make UnitRegistry.__contains__ a bit smarter

docs/usage.rst

@@ -732,11 +732,16 @@ Sometimes it is convenient to create a unit registry containing new units that a
   >>> reg = UnitRegistry()
   >>> reg.add("code_length", base_value=10.0, dimensions=length,
   ...         tex_repr=r"\rm{Code Length}")
+  >>> 'code_length' in reg
+  True
   >>> u = Unit('code_length', registry=reg)
   >>> data = 3*u
   >>> print(data)
   3.0 code_length
 
+As you can see, you can test whether a unit name is in a registry using the
+Python ``in`` operator.
+
 In an application that depends on ``unyt``, it is often convenient to define
 methods or functions to automatically attach the correct unit registry to a set
 unit object. For example, consider a ``Simulation`` class. Let's give this class

unyt/tests/test_unit_registry.py

@@ -66,3 +66,13 @@ def test_prefixable_units():
     mfoobar = Unit('mfoobar', registry=ureg)
     foobar = Unit('foobar', registry=ureg)
     assert (1*foobar)/(1*mfoobar) == 1000
+
+
+def test_registry_contains():
+    ureg = UnitRegistry()
+    assert 'm' in ureg
+    assert 'cm' in ureg
+    assert 'erg' in ureg
+    assert 'Merg' in ureg
+    assert 'foobar' not in ureg
+    assert Unit('m', registry=ureg) in ureg

unyt/unit_object.py

@@ -64,14 +64,11 @@
     UnitsNotReducible,
 )
 from unyt._physical_ratios import speed_of_light_cm_per_s
-from unyt._unit_lookup_table import (
-    unit_prefixes,
-    latex_prefixes,
-)
 from unyt.unit_registry import (
     default_unit_registry,
+    _lookup_unit_symbol,
     UnitRegistry,
-    UnitParseError
+    UnitParseError,
 )
 
 sympy_one = sympify(1)
@@ -898,85 +895,6 @@ def _get_unit_data_from_expr(unit_expr, unit_symbol_lut):
                          "objects." % str(unit_expr))
 
 
-def _lookup_unit_symbol(symbol_str, unit_symbol_lut):
-    """
-    Searches for the unit data tuple corresponding to the given symbol.
-
-    Parameters
-    ----------
-    symbol_str : str
-        The unit symbol to look up.
-    unit_symbol_lut : dict
-        Dictionary with symbols as keys and unit data tuples as values.
-
-    """
-    if symbol_str in unit_symbol_lut:
-        # lookup successful, return the tuple directly
-        return unit_symbol_lut[symbol_str]
-
-    # could still be a known symbol with a prefix
-    possible_prefix = symbol_str[0]
-
-    if symbol_str[:2] == 'da':
-        possible_prefix = 'da'
-
-    if possible_prefix in unit_prefixes:
-        # the first character could be a prefix, check the rest of the symbol
-        symbol_wo_pref = symbol_str[1:]
-
-        # deca is the only prefix with length 2
-        if symbol_str[:2] == 'da':
-            symbol_wo_pref = symbol_str[2:]
-            possible_prefix = 'da'
-
-        prefixable_units = [u for u in unit_symbol_lut
-                            if unit_symbol_lut[u][4]]
-
-        unit_is_si_prefixable = (symbol_wo_pref in unit_symbol_lut and
-                                 symbol_wo_pref in prefixable_units)
-
-        if unit_is_si_prefixable is True:
-            # lookup successful, it's a symbol with a prefix
-            unit_data = unit_symbol_lut[symbol_wo_pref]
-            prefix_value = unit_prefixes[possible_prefix]
-
-            if possible_prefix in latex_prefixes:
-                latex_repr = symbol_str.replace(
-                    possible_prefix, '{'+latex_prefixes[possible_prefix]+'}')
-            else:
-                # Need to add some special handling for comoving units
-                # this is fine for now, but it wouldn't work for a general
-                # unit that has an arbitrary LaTeX representation
-                if symbol_wo_pref != 'cm' and symbol_wo_pref.endswith('cm'):
-                    sub_symbol_wo_prefix = symbol_wo_pref[:-2]
-                    sub_symbol_str = symbol_str[:-2]
-                else:
-                    sub_symbol_wo_prefix = symbol_wo_pref
-                    sub_symbol_str = symbol_str
-                latex_repr = unit_data[3].replace(
-                    '{' + sub_symbol_wo_prefix + '}',
-                    '{' + sub_symbol_str + '}')
-
-            # Leave offset and dimensions the same, but adjust scale factor and
-            # LaTeX representation
-            ret = (unit_data[0] * prefix_value, unit_data[1], unit_data[2],
-                   latex_repr, False)
-
-            unit_symbol_lut[symbol_str] = ret
-
-            return ret
-
-    # no dice
-    if symbol_str.startswith('code_'):
-        raise UnitParseError(
-            "Code units have not been defined. \n"
-            "Try creating the array or quantity using ds.arr or ds.quan "
-            "instead.")
-    else:
-        raise UnitParseError("Could not find unit symbol '%s' in the provided "
-                             "symbols." % symbol_str)
-
-
 def _validate_dimensions(dimensions):
     if isinstance(dimensions, Mul):
         for dim in dimensions.args:

unyt/unit_registry.py

@@ -19,7 +19,11 @@
     SymbolNotFoundError,
     UnitParseError,
 )
-from unyt._unit_lookup_table import default_unit_symbol_lut
+from unyt._unit_lookup_table import (
+    default_unit_symbol_lut,
+    unit_prefixes,
+    latex_prefixes,
+)
 from hashlib import md5
 import six
 from sympy import (
@@ -46,7 +50,13 @@ def __getitem__(self, key):
         return self.lut[key]
 
     def __contains__(self, item):
-        return item in self.lut
+        if str(item) in self.lut:
+            return True
+        try:
+            _lookup_unit_symbol(str(item), self.lut)
+            return True
+        except UnitParseError:
+            return False
 
     def pop(self, item):
         if item in self.lut:
@@ -230,5 +240,78 @@ def list_same_dimensions(self, unit_object):
         return equiv
 
 
+def _lookup_unit_symbol(symbol_str, unit_symbol_lut):
+    """
+    Searches for the unit data tuple corresponding to the given symbol.
+
+    Parameters
+    ----------
+    symbol_str : str
+        The unit symbol to look up.
+    unit_symbol_lut : dict
+        Dictionary with symbols as keys and unit data tuples as values.
+
+    """
+    if symbol_str in unit_symbol_lut:
+        # lookup successful, return the tuple directly
+        return unit_symbol_lut[symbol_str]
+
+    # could still be a known symbol with a prefix
+    possible_prefix = symbol_str[0]
+
+    if symbol_str[:2] == 'da':
+        possible_prefix = 'da'
+
+    if possible_prefix in unit_prefixes:
+        # the first character could be a prefix, check the rest of the symbol
+        symbol_wo_pref = symbol_str[1:]
+
+        # deca is the only prefix with length 2
+        if symbol_str[:2] == 'da':
+            symbol_wo_pref = symbol_str[2:]
+            possible_prefix = 'da'
+
+        prefixable_units = [u for u in unit_symbol_lut
+                            if unit_symbol_lut[u][4]]
+
+        unit_is_si_prefixable = (symbol_wo_pref in unit_symbol_lut and
+                                 symbol_wo_pref in prefixable_units)
+
+        if unit_is_si_prefixable is True:
+            # lookup successful, it's a symbol with a prefix
+            unit_data = unit_symbol_lut[symbol_wo_pref]
+            prefix_value = unit_prefixes[possible_prefix]
+
+            if possible_prefix in latex_prefixes:
+                latex_repr = symbol_str.replace(
+                    possible_prefix, '{'+latex_prefixes[possible_prefix]+'}')
+            else:
+                # Need to add some special handling for comoving units
+                # this is fine for now, but it wouldn't work for a general
+                # unit that has an arbitrary LaTeX representation
+                if symbol_wo_pref != 'cm' and symbol_wo_pref.endswith('cm'):
+                    sub_symbol_wo_prefix = symbol_wo_pref[:-2]
+                    sub_symbol_str = symbol_str[:-2]
+                else:
+                    sub_symbol_wo_prefix = symbol_wo_pref
+                    sub_symbol_str = symbol_str
+                latex_repr = unit_data[3].replace(
+                    '{' + sub_symbol_wo_prefix + '}',
+                    '{' + sub_symbol_str + '}')
+
+            # Leave offset and dimensions the same, but adjust scale factor and
+            # LaTeX representation
+            ret = (unit_data[0] * prefix_value, unit_data[1], unit_data[2],
+                   latex_repr, False)
+
+            unit_symbol_lut[symbol_str] = ret
+
+            return ret
+
+    # no dice
+    raise UnitParseError("Could not find unit symbol '%s' in the provided "
+                         "symbols." % symbol_str)
+
+
 #: The default unit registry
 default_unit_registry = UnitRegistry()