Add comments and improve docstrings

2 files changed, 18 insertions, 2 deletions

diff --git a/mdpoly/expressions.py b/mdpoly/expressions.py
index 9fa07f2..159769e 100644
--- a/mdpoly/expressions.py
+++ b/mdpoly/expressions.py
@@ -164,6 +164,10 @@ class WithOps:
         return self.unwrap()
 
     def __exit__(self, *ex):
+        """ Context manager exit.
+
+        See :py:meth:`mdpoly.expresssions.WithOps.__enter__`.
+        """
         return False # Propagate exceptions
 
     def __str__(self):
@@ -217,9 +221,19 @@ class WithOps:
         return wrapper
 
     @staticmethod
-    def wrap_result(meth: Callable[[WithOps, Any], Expr]) -> Callable[[WithOps, WithOps], WithOps]:
+    def wrap_result(meth: Callable[[WithOps, Any], Expr]) -> Callable[[WithOps, Any], WithOps]:
+        """ Take a method and wrap its result type.
+
+        Turns method `(WithOps, Any) -> Expr)` into 
+        `(WithOps, Any) -> WithOps)`. method arguments are left unchanged.
+        This is only for conveniente to avoid always having to wrap the result
+        by hand.
+        """
         @wraps(meth)
         def meth_wrapper(self, *args, **kwargs) -> WithOps:
+            # Why type(self)? Because if we are wrapping a subtype of WithOps,
+            # eg. from an extension we want to preserve its type. See for
+            # example mdpoly.test.TestExtensions.
             return type(self)(expr=meth(self, *args, **kwargs))
         return meth_wrapper
 
@@ -227,6 +241,7 @@ class WithOps:
 
     @classmethod
     def _ensure_is_withops(cls, obj: WithOps | Expr | Number) -> WithOps:
+        """ Ensures that the given object is wrapped with type WithOps. """
         if isinstance(obj, WithOps):
             return obj
 
diff --git a/mdpoly/operations/__init__.py b/mdpoly/operations/__init__.py
index 752358c..c75608d 100644
--- a/mdpoly/operations/__init__.py
+++ b/mdpoly/operations/__init__.py
@@ -61,7 +61,8 @@ class Reducible(Expr):
 
     @abstractmethod
     def reduce(self) -> Expr:
-        """ Reduce the expression to its basic elements """
+        """ Reduce to an expression made of other operations that have already
+        implemented :py:meth:`mdpoly.abc.Expr.to_repr`. """
 
 
 # TODO: review this idea before implementing