code: def unitized(unit):
      
      def wrap(f):
          def wrapped_f(*args, **kwargs):
              val = f(*args, **kwargs)
              unit_type = _UNAME2UTYPE[unit]
  
              if isinstance(val, FloatWithUnit) or isinstance(val, ArrayWithUnit):
                  return val.to(unit)
  
              elif isinstance(val, collections.abc.Sequence):
                  # TODO: why don't we return a ArrayWithUnit?
                  # This complicated way is to ensure the sequence type is
                  # preserved (list or tuple).
                  return val.__class__([FloatWithUnit(i, unit_type=unit_type,
                                                      unit=unit) for i in val])
              elif isinstance(val, collections.Mapping):
                  for k, v in val.items():
                      val[k] = FloatWithUnit(v, unit_type=unit_type, unit=unit)
              elif isinstance(val, numbers.Number):
                  return FloatWithUnit(val, unit_type=unit_type, unit=unit)
              elif val is None:
                  pass
              else:
                  raise TypeError("Don't know how to assign units to %s" % str(val))
              return val
          return wrapped_f
      return wrap
docs: Useful decorator to assign units to the output of a function. You can also
      use it to standardize the output units of a function that already returns
      a FloatWithUnit or ArrayWithUnit. For sequences, all values in the sequences
      are assigned the same unit. It works with Python sequences only. The creation
      of numpy arrays loses all unit information. For mapping types, the values
      are assigned units.
  
      Args:
          unit: Specific unit (eV, Ha, m, ang, etc.).
  
      Example usage::
  
          @unitized(unit="kg")
          def get_mass():
              return 123.45