Merge pull request #3464 from Robbybp/cuid-from-cuid

Allow construction of CUID from another CUID

Author: blnicho

pyomo/contrib/mpc/data/get_cuid.py

@@ -37,11 +37,8 @@ def get_indexed_cuid(var, sets=None, dereference=None, context=None):
         ComponentUID corresponding to the provided ``var`` and sets
 
     """
-    # TODO: Does this function have a good name?
     # Should this function be generalized beyond a single indexing set?
-    if isinstance(var, ComponentUID):
-        return var
-    elif isinstance(var, (str, IndexedComponent_slice)):
+    if isinstance(var, (str, IndexedComponent_slice, ComponentUID)):
         # TODO: Raise error if string and context is None
         return ComponentUID(var, context=context)
     # At this point we are assuming var is a Pyomo Var or VarData object.

pyomo/core/base/block.py

@@ -921,11 +921,7 @@ def find_component(self, label_or_component):
             a matching component is not found, None is returned.
 
         """
-        if type(label_or_component) is ComponentUID:
-            cuid = label_or_component
-        else:
-            cuid = ComponentUID(label_or_component)
-        return cuid.find_component_on(self)
+        return ComponentUID(label_or_component).find_component_on(self)
 
     @contextmanager
     def _declare_reserved_components(self):

pyomo/core/base/componentuid.py

@@ -43,6 +43,12 @@ def _index_repr(x):
     return __index_repr(x, _pickle)
 
 
+def _context_err(_type):
+    raise ValueError(
+        f"Context is not allowed when initializing a ComponentUID from {_type}."
+    )
+
+
 class ComponentUID(object):
     """
     A Component unique identifier
@@ -78,15 +84,15 @@ def __init__(self, component, cuid_buffer=None, context=None):
         # the string representation.
         if isinstance(component, str):
             if context is not None:
-                raise ValueError(
-                    "Context is not allowed when initializing a "
-                    "ComponentUID object from a string type"
-                )
+                _context_err(str)
             try:
                 self._cids = tuple(self._parse_cuid_v2(component))
             except (OSError, IOError):
                 self._cids = tuple(self._parse_cuid_v1(component))
-
+        elif type(component) is ComponentUID:
+            if context is not None:
+                _context_err(ComponentUID)
+            self._cids = component._cids
         elif type(component) is IndexedComponent_slice:
             self._cids = tuple(
                 self._generate_cuid_from_slice(component, context=context)

pyomo/core/tests/unit/test_componentuid.py

@@ -81,11 +81,7 @@ def test_genFromComponent_context(self):
             ValueError, r"Context 'b\[1,'2'\]' does not apply to component 's'"
         ):
             ComponentUID(self.m.s, context=self.m.b[1, '2'])
-        with self.assertRaisesRegex(
-            ValueError,
-            "Context is not allowed when initializing a ComponentUID "
-            "object from a string type",
-        ):
+        with self.assertRaisesRegex(ValueError, "Context is not allowed"):
             ComponentUID("b[1,2].c.a[2]", context=self.m.b[1, '2'])
 
     def test_parseFromString(self):
@@ -1248,6 +1244,26 @@ def test_cuid_from_slice_errors(self):
         ):
             cuid = ComponentUID(_slice)
 
+    def test_cuid_from_cuid(self):
+        def assert_equal(cuid1, cuid2):
+            self.assertEqual(cuid1, cuid2)
+            self.assertFalse(cuid1 is cuid2)
+
+        cuid_str = ComponentUID("b.var[1]")
+        cuid_str_2 = ComponentUID(cuid_str)
+        assert_equal(cuid_str, cuid_str_2)
+
+        cuid_comp = ComponentUID(self.m.b[1, 1].c)
+        cuid_comp_2 = ComponentUID(cuid_comp)
+        assert_equal(cuid_str, cuid_str_2)
+
+        cuid_slice = ComponentUID(self.m.b[1, :].c)
+        cuid_slice_2 = ComponentUID(cuid_slice)
+        assert_equal(cuid_slice, cuid_slice_2)
+
+        with self.assertRaisesRegex(ValueError, "Context is not allowed"):
+            ComponentUID(cuid_comp, context=self.m.b[1, 1])
+
 
 if __name__ == "__main__":
     unittest.main()