Make PolyMatrixDict accept indexing p[row, col]

1 files changed, 11 insertions, 9 deletions

diff --git a/polymatrix/polymatrix/typing.py b/polymatrix/polymatrix/typing.py
index 43d7210..1de5c36 100644
--- a/polymatrix/polymatrix/typing.py
+++ b/polymatrix/polymatrix/typing.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 from typing import NamedTuple, Iterable
+from collections import UserDict
 
 # TODO: remove these types, they are here for backward compatiblity
 MonomialData = tuple[tuple[int, int], ...]
@@ -34,17 +35,10 @@ class MonomialIndex(tuple[VariableIndex]):
 class PolyDict(dict[MonomialIndex, int | float]):
     """ Polynomial, stored as a dictionary. """
 
-    # NOTE: this class directly inherits from dict instead of
-    # collections.UserDict, hence it is not possible to override any of the
-    # __dunder__ methods. If you like to change behaviour of one of these
-    # methods change to inherit from UserDict, this could incur a slight
-    # performance cost, because dict is straight CPython and thus faster
-
     @staticmethod
     def empty() -> PolyDict:
         return PolyDict({})
 
-
     def terms(self) -> Iterable[tuple[MonomialIndex, int | float]]:
         """ Iterate over terms with a non-zero coefficient. """
         # This is an alias for readability
@@ -67,7 +61,15 @@ class MatrixIndex(NamedTuple):
     col: int
 
 
-class PolyMatrixDict(dict[MatrixIndex, PolyDict]):
+class PolyMatrixDict(UserDict[MatrixIndex, PolyDict]):
     """ Matrix whose entries are polynomials, stored as a dictionary. """
 
-    # NOTE: Same as NOTE in PolyDict.
+    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):
+        if not isinstance(key, MatrixIndex):
+            key = MatrixIndex(*key)
+        return super().__setitem__(key, value)