Make State.from_index accept PolyVarIndex indices

2 files changed, 7 insertions, 2 deletions

diff --git a/README.rst b/README.rst
index f4deb4d..c539b15 100644
--- a/README.rst
+++ b/README.rst
@@ -36,7 +36,7 @@ to give an overview.
             for term in sparse.terms(entry):
                 monomial_str = ""
                 for idx in term:
-                    var = state.from_index(idx.var_idx)
+                    var = state.from_index(idx)
                     monomial_str += f"{var.name}^{idx.power} "
 
                 # Get the coefficient
diff --git a/mdpoly/state.py b/mdpoly/state.py
index f22d502..70fca24 100644
--- a/mdpoly/state.py
+++ b/mdpoly/state.py
@@ -2,6 +2,8 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 from dataclasses import dataclass, field
 
+from .types import PolyVarIndex
+
 if TYPE_CHECKING:
     from .types import Number
     from .leaves import Var, Param, Const
@@ -33,10 +35,13 @@ class State:
 
         return self.variables[var]
 
-    def from_index(self, index: Index) -> Var | Const:
+    def from_index(self, index: Index | PolyVarIndex) -> Var | Const:
         """ Get a variable object from the index.
         This is a reverse lookup operation and it is **not**
         :math:`\mathcal{O}(1)`!"""
+        if isinstance(index, PolyVarIndex):
+            index = index.var_idx
+
         if index == -1:
             from .leaves import Const
             return Const(value=1, name="1")