Add helper methods to index types, make indexing more flexible.

Following robustness principle: be conservative in what you do, be
liberal in what you accept from others

1 files changed, 28 insertions, 4 deletions

diff --git a/polymatrix/polymatrix/typing.py b/polymatrix/polymatrix/typing.py
index 1de5c36..4baff5a 100644
--- a/polymatrix/polymatrix/typing.py
+++ b/polymatrix/polymatrix/typing.py
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import NamedTuple, Iterable
+from typing import NamedTuple, Iterable, cast
 from collections import UserDict
 
 # TODO: remove these types, they are here for backward compatiblity
@@ -15,11 +15,14 @@ class VariableIndex(NamedTuple):
     variable: int # index in ExpressionState object
     power: int
 
-
     def __lt__(self, other):
         """ Variables indices can be sorted with respect to variable index. """
         return self.variable < other.variable
 
+    @staticmethod
+    def constant() -> VariableIndex:
+        return cast(VariableIndex, tuple())
+
 
 class MonomialIndex(tuple[VariableIndex]):
     """
@@ -30,15 +33,28 @@ class MonomialIndex(tuple[VariableIndex]):
     def empty() -> MonomialIndex:
         return MonomialIndex(tuple())
 
+    @staticmethod
+    def constant() -> MonomialIndex:
+        return MonomialIndex((VariableIndex.constant(),))
 
 
-class PolyDict(dict[MonomialIndex, int | float]):
+class PolyDict(UserDict[MonomialIndex, int | float]):
     """ Polynomial, stored as a dictionary. """
 
     @staticmethod
     def empty() -> PolyDict:
         return PolyDict({})
 
+    def __getitem__(self, key: Iterable[VariableIndex] | MonomialIndex) -> int | float:
+        if not isinstance(key, MonomialIndex):
+            key = MonomialIndex(key)
+        return super().__getitem__(key)
+
+    def __setitem__(self, key: Iterable[VariableIndex] | MonomialIndex, value: int | float):
+        if not isinstance(key, MonomialIndex):
+            key = MonomialIndex(key)
+        return super().__setitem__(key, value)
+
     def terms(self) -> Iterable[tuple[MonomialIndex, int | float]]:
         """ Iterate over terms with a non-zero coefficient. """
         # This is an alias for readability
@@ -64,12 +80,20 @@ class MatrixIndex(NamedTuple):
 class PolyMatrixDict(UserDict[MatrixIndex, PolyDict]):
     """ Matrix whose entries are polynomials, stored as a dictionary. """
 
+    @staticmethod
+    def empty() -> PolyMatrixDict:
+        return PolyMatrixDict({})
+
     def __getitem__(self, key: tuple[int, int] | MatrixIndex) -> PolyDict:
         if not isinstance(key, MatrixIndex):
             key = MatrixIndex(*key)
         return super().__getitem__(key)
 
-    def __setitem__(self, key: tuple[int, int] | MatrixIndex, value: PolyDict):
+    def __setitem__(self, key: tuple[int, int] | MatrixIndex, value: dict | PolyDict):
         if not isinstance(key, MatrixIndex):
             key = MatrixIndex(*key)
+
+        if not isinstance(value, PolyDict):
+            value = PolyDict(value)
+
         return super().__setitem__(key, value)