some documentation

Signed-off-by: Nathaniel Starkman (@nstarman) <nstarkman@protonmail.com>

astropy/units/quantity_helper/helpers.py

@@ -26,10 +26,6 @@ def _d(unit):
         return unit
 
 
-def _I(x):  # identity function
-    return x
-
-
 def get_converter(from_unit, to_unit):
     """Like Unit._get_converter, except returns None if no scaling is needed,
     i.e., if the inferred scale is unity."""
@@ -325,15 +321,50 @@ def helper_clip(f, unit1, unit2, unit3):
 
 
 # HELPER NARGS
-def register_ufunc(ufunc, nin, nout, inunits, outunits):
-    from astropy.units import UnitBase, FunctionUnitBase
-
-    if isinstance(inunits, (UnitBase, FunctionUnitBase)):  # scalar -> list
-        inunits = [inunits]
+def register_ufunc(ufunc, nin, nout, inunits, ounits):
+    """
+    Register `~numpy.ufunc` in ``UFUNC_HELPERS``, along with the conversion
+    functions necessary to strip input units and assign output units. ufuncs
+    operate on only recognized `~numpy.dtype`s (e.g. int, float32), so units
+    must be removed beforehand and replaced afterwards. Therefore units MUST
+    BE KNOWN a priori, as they will not be propagated by the astropy machinery.
+
+    Parameters
+    ----------
+    ufunc : `~numpy.ufunc`
+    nin, nout : int
+        Number of ufunc's inputs and outputs
+    inunits, ounits : sequence[unit-like]
+        Sequence of the input and output units, respectively.
+
+    """
+    from astropy.units import Unit
+
+    # process sequence[unit-like] -> sequence[unit]
+    inunits = [Unit(iu) for iu in inunits]
+    ounits = [Unit(ou) for ou in ounits]
 
     def helper_nargs(f, *units):
-        converters = [get_converter(frm or to, to) for frm, to in zip(units, inunits)]
-        return converters, outunits
+        """Helper function to convert input units and assign output units.
+
+        Parameters
+        ----------
+        f : callable
+        *units : `~astropy.units.UnitBase` or None
+            The units of the inputs. If None (a unitless input) the unit is
+            assumed to be the one specified in
+            `~astropy.units.quantity_helper.helpers.register_ufunc`
+
+        Returns
+        -------
+        converters : sequence[callable]
+        ounits : sequence[unit-like]
+
+        """
+        # no units assumed to be in inunits
+        converters = [get_converter(frm or to, to)
+                      for frm, to in zip(units, inunits)]
+        return converters, ounits
 
     UFUNC_HELPERS[ufunc] = helper_nargs
 

astropy/units/utils.py

@@ -299,39 +299,82 @@ def quantity_asanyarray(a, dtype=None):
 
 # ------------------------------------------------------------------------------
 
-def quantity_frompyfunc(func, nin, nout, inunits=None, ounits=None, *, identity=None):
-    # This function is dangerous b/c ounits is the final word on
-    # output units
-    from astropy.units import UnitBase, FunctionUnitBase
+def _is_seq_ulike(annotation):
+    from astropy.units import UnitBase
 
-    sig = inspect.signature(func)
+    is_unit = isinstance(annotation, UnitBase)
+    is_unit_sequence = (isinstance(annotation, Sequence) and
+                        all(isinstance(x, UnitBase) for x in annotation))
+    return True if (is_unit or is_unit_sequence) else False
 
-    def is_unitlike(annotation):
-        is_unit = isinstance(annotation, (UnitBase, FunctionUnitBase))
-        is_unit_sequence = (
-            isinstance(annotation, Sequence) and
-            all(isinstance(x, (UnitBase, FunctionUnitBase)) for x in annotation))
-        if is_unit or is_unit_sequence:
-            return True
 
-    if inunits is None:
-        svals = tuple(sig.parameters.values())
+def quantity_frompyfunc(func, nin, nout, inunits=None, ounits=None,
+                        *, identity=None):
+    """Quantity-aware `~numpy.frompyfunc`.
+
+    `~numpy.ufunc`s operate on only recognized `~numpy.dtype`s (e.g. float32),
+    so units must be removed beforehand and replaced afterwards. Therefore units
+    MUST BE KNOWN a priori, as they will not be propagated by the astropy
+    machinery.
+
+    Parameters
+    ----------
+    func : callable
+    nin, nout : int
+        Number of ufunc's inputs and outputs
+    inunits, ounits : sequence[unit-like] (optional)
+        Sequence of the input and output units, respectively.
+
+        .. warning::
+            Inputs will be converted to these units before being passed to the
+            returned `~numpy.ufunc`. Outputs will be assigned these units.
+            Make sure these are the correct units.
+
+    identity : object (optional, keyword-only)
+        The value to use for the `~numpy.ufunc.identity` attribute of the
+        resulting object. If specified, this is equivalent to setting the
+        underlying C ``identity`` field to ``PyUFunc_IdentityValue``.
+        If omitted, the identity is set to ``PyUFunc_None``. Note that this is
+        _not_ equivalent to setting the identity to ``None``, which implies the
+        operation is reorderable.
+
+    Returns
+    -------
+    `~numpy.ufunc`
 
-        inunits = [p.annotation if is_unitlike(p.annotation) else None
+    """
+    from astropy.units import Unit
+
+    # -------------------------
+    # determine units by introspection
+
+    sig = inspect.signature(func)
+
+    # input units
+    if inunits is not None:
+        inunits = [Unit(iu) for iu in inunits]  # seq[unit-like] -> seq[unit]
+    else:
+        svals = tuple(sig.parameters.values())
+        # TODO! more robust. what if no annotations?
+        inunits = [Unit(p.annotation) if _is_seq_ulike(p.annotation) else None
                    for p in svals]
 
-    if ounits is None:
+    # output units
+    if ounits is not None:
+        ounits = [Unit(ou) for ou in ounits]
+    else:
         ra = sig.return_annotation
-        # TODO! more intelligent detection of unit returns
-        if is_unitlike(ra):
+        if _is_seq_ulike(ra):
             ounits = ra
 
+    # -------------------------
     # make ufunc
+
     ufunc = np.frompyfunc(func, nin, nout, identity=identity)
 
-    # register ufunc with Quantity
+    # register ufunc with Quantity machinery
     from astropy.units.quantity_helper.helpers import register_ufunc
-    register_ufunc(ufunc, nin, nout, inunits, ounits)
+    register_ufunc(ufunc, nin=nin, nout=nout, inunits=inunits, ounits=ounits)
 
     return ufunc