docs: print docstr before returning in MathPadConstructor._repr_latex_()

mathpad/_quality_of_life.py

@@ -1,11 +1,16 @@
 import sympy
 import inspect
-from typing import Callable, cast, TypeVar
+from typing import TYPE_CHECKING, Callable, Union, Generic, TypeVar
 
 from mathpad.units import seconds, radians
 from mathpad.val import Dimensionless, Val
 from mathpad.dimensions import Angle
 
+if TYPE_CHECKING:
+    from mathpad.equation import Equation
+    from mathpad.vector import Vector
+    from mathpad.matrix import Matrix
+
 # from mathpad.global_options import _global_options
 
 # don't display this symbol on definition
@@ -18,36 +23,57 @@
 e = Dimensionless(1, sympy.E) # type: ignore
 dimensionless = Dimensionless(1, 1)
 
+MathPadObject = Union[Val, 'Equation', 'Vector', 'Matrix']
+T = TypeVar("T", bound=MathPadObject)
 
-
-class MathPadConstructor:
+class MathPadConstructor(Generic[T]):
     """
     A function decorator that adds a _repr_latex_ method to the function so that
     it can be displayed in IPython.
     """
 
-    def __init__(self, fn):
+    def __init__(self, fn: Callable[..., T]):
         self.fn = fn
 
     def _repr_latex_(self):
-        "Calls the function with created symbolic values"
+        "Calls the function with created symbolic values to produce a LaTeX representation"
 
         argspec = inspect.signature(self.fn)
 
-        args = {
+        arg2dimension = {
             # produce a symbolic value for each argument, with base units of that argument's dimension
-            param.name: param.name * Val(param.annotation.__args__[0].base_units)
+            param.name: param.annotation.__args__[0]
             for param in argspec.parameters.values()
         }
+
+        args = {
+            # produce a symbolic value for each argument, with base units of that argument's dimension
+            pname: pname * Val(dimension.base_units)
+            for pname, dimension in arg2dimension.items()
+        }
+
+        argspec_str = ", ".join(f"{pname}: {dimension.__name__}" for pname, dimension in arg2dimension.items())
+
+        # TODO: parse the type signature to determine what parameters are kwargs only rather than assume all are
+        print(f"{self.fn.__name__}(*, {argspec_str})")
+
+        # print the docstr if we have one
+        if self.fn.__doc__:
+            # hide the example section if it exists
+            pre_example_docstr, *_ = self.fn.__doc__.split("Example:")
+            print(pre_example_docstr.rstrip())
+
+        mp_obj = self.fn(**args)
 
-        return self.fn(**args)._repr_latex_()
+        return mp_obj._repr_latex_()
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs) -> T:
         return self.fn(*args, **kwargs)
 
 
 F = TypeVar("F", bound=Callable)
 
+# lie about the return type so that the function's signature is preserved
 def mathpad_constructor(fn: F) -> F:
     "function decorator that adds a _repr_latex_ method to the function so that it can be displayed in IPython"
-    return cast(type(fn), MathPadConstructor(fn))
+    return MathPadConstructor(fn) # type: ignore