Merge pull request #338 from candleindark/add-class-improv

BF: Fix bug and other improvement in `SchemaBuilder.add_class()`

Author: cmungall

linkml_runtime/utils/schema_builder.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass, fields
 from typing import Dict, List, Union, Optional
 
 from linkml_runtime.linkml_model import (ClassDefinition, EnumDefinition,
@@ -16,7 +16,7 @@ class SchemaBuilder:
 
     Example:
 
-        >>> from linkml.utils.schema_builder import SchemaBuilder
+        >>> from linkml_runtime.utils.schema_builder import SchemaBuilder
         >>> sb = SchemaBuilder('test-schema')
         >>> sb.add_class('Person', slots=['name', 'age'])
         >>> sb.add_class('Organization', slots=['name', 'employees'])
@@ -57,26 +57,47 @@ def add_class(
         cls: Union[ClassDefinition, Dict, str],
         slots: List[Union[str, SlotDefinition]] = None,
         slot_usage: Dict[str, SlotDefinition] = None,
-        replace_if_present=False,
-        use_attributes=False,
+        replace_if_present: bool = False,
+        use_attributes: bool = False,
         **kwargs,
     ) -> "SchemaBuilder":
         """
         Adds a class to the schema.
 
         :param cls: name, dict object, or ClassDefinition object to add
-        :param slots: slot of slot names or slot objects.
-        :param slot_usage: slots keyed by slot name
+        :param slots: list of slot names or slot objects. This must be a list of
+            `SlotDefinition` objects if `use_attributes=True`
+        :param slot_usage: slots keyed by slot name (ignored if `use_attributes=True`)
         :param replace_if_present: if True, replace existing class if present
+        :param use_attributes: Whether to specify the given slots as an inline
+            definition of slots, attributes, in the class definition
         :param kwargs: additional ClassDefinition properties
         :return: builder
         :raises ValueError: if class already exists and replace_if_present=False
         """
+        if slots is None:
+            slots = []
+        if slot_usage is None:
+            slot_usage = {}
+
         if isinstance(cls, str):
             cls = ClassDefinition(cls, **kwargs)
-        if isinstance(cls, dict):
+        elif isinstance(cls, dict):
             cls = ClassDefinition(**{**cls, **kwargs})
-        if cls.name is self.schema.classes and not replace_if_present:
+        else:
+            # Ensure that `cls` is a `ClassDefinition` object
+            if not isinstance(cls, ClassDefinition):
+                msg = (
+                    f"cls must be a string, dict, or ClassDefinition, "
+                    f"not {type(cls)!r}"
+                )
+                raise TypeError(msg)
+
+            cls_as_dict = {f.name: getattr(cls, f.name) for f in fields(cls)}
+
+            cls = ClassDefinition(**{**cls_as_dict, **kwargs})
+
+        if cls.name in self.schema.classes and not replace_if_present:
             raise ValueError(f"Class {cls.name} already exists")
         self.schema.classes[cls.name] = cls
         if use_attributes:
@@ -88,18 +109,14 @@ def add_class(
                         f"If use_attributes=True then slots must be SlotDefinitions"
                     )
         else:
-            if slots is not None:
-                for s in slots:
-                    cls.slots.append(s.name if isinstance(s, SlotDefinition) else s)
-                    if isinstance(s, str) and s in self.schema.slots:
-                        # top-level slot already exists
-                        continue
-                    self.add_slot(s, replace_if_present=replace_if_present)
-            if slot_usage:
-                for k, v in slot_usage.items():
-                    cls.slot_usage[k] = v
-        for k, v in kwargs.items():
-            setattr(cls, k, v)
+            for s in slots:
+                cls.slots.append(s.name if isinstance(s, SlotDefinition) else s)
+                if isinstance(s, str) and s in self.schema.slots:
+                    # top-level slot already exists
+                    continue
+                self.add_slot(s, replace_if_present=replace_if_present)
+            for k, v in slot_usage.items():
+                cls.slot_usage[k] = v
         return self
 
     def add_slot(

tests/test_utils/test_schema_builder.py

@@ -0,0 +1,146 @@
+from typing import Optional, List, Union, Any, Dict
+from dataclasses import fields
+
+import pytest
+
+from linkml_runtime.utils.schema_builder import SchemaBuilder
+from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
+
+
+@pytest.mark.parametrize("replace_if_present", [True, False])
+def test_add_existing_class(replace_if_present):
+    """
+    Test the case of adding a class with a name that is the same as a class that already
+    exists in the schema
+    """
+    builder = SchemaBuilder()
+
+    cls_name = "Person"
+
+    # Add a class to the schema
+    cls = ClassDefinition(name=cls_name, slots=["name"])
+    builder.add_class(cls)
+
+    # Add another class with the same name to the schema
+    cls_repeat = ClassDefinition(name=cls_name, slots=["age"])
+
+    if replace_if_present:
+        builder.add_class(cls_repeat, replace_if_present=True)
+        assert builder.schema.classes[cls_name].slots == ["age"]
+    else:
+        with pytest.raises(ValueError, match=f"Class {cls_name} already exists"):
+            builder.add_class(cls_repeat)
+        assert builder.schema.classes[cls_name].slots == ["name"]
+
+
+@pytest.mark.parametrize(
+    "slots",
+    [
+        None,
+        ["name", "age"],
+        ["name", SlotDefinition(name="age")],
+        [SlotDefinition(name="name"), SlotDefinition(name="age")],
+    ],
+)
+@pytest.mark.parametrize("use_attributes", [True, False])
+def test_add_class_with_slot_additions(
+    slots: Optional[List[Union[str, SlotDefinition]]], use_attributes: bool
+):
+    """
+    Test adding a class with separate additional slots specification
+    """
+    # If `slots` is None, it should be treated as if it were an empty list
+    if slots is None:
+        slots = []
+
+    slot_names = {s if isinstance(s, str) else s.name for s in slots}
+
+    builder = SchemaBuilder()
+
+    cls_name = "Person"
+
+    # Add a class to the schema
+    cls = ClassDefinition(name=cls_name)
+
+    if use_attributes:
+        # === The case where the slots specified should be added to the `attributes`
+        # meta slot of the class ===
+        if any(not isinstance(s, SlotDefinition) for s in slots):
+            with pytest.raises(
+                ValueError,
+                match="If use_attributes=True then slots must be SlotDefinitions",
+            ):
+                builder.add_class(cls, slots=slots, use_attributes=use_attributes)
+        else:
+            builder.add_class(cls, slots=slots, use_attributes=use_attributes)
+            assert builder.schema.classes[cls_name].attributes.keys() == slot_names
+    else:
+        # === The case where the slots specified should be added to the `slots`
+        # meta slot of the schema ===
+        builder.add_class(cls, slots=slots, use_attributes=use_attributes)
+        assert builder.schema.slots.keys() == slot_names
+
+
+@pytest.mark.parametrize(
+    ("cls", "extra_kwargs", "expected_added_class"),
+    [
+        ("Person", {}, ClassDefinition(name="Person")),
+        (2, {}, None),  # Invalid type for `cls`
+        ("Person", {"tree_root": True}, ClassDefinition(name="Person", tree_root=True)),
+        ("Person", {"ijk": True}, None),  # Invalid extra kwarg
+        (
+            {"name": "Person", "tree_root": False},
+            {"tree_root": True},
+            ClassDefinition(name="Person", tree_root=True),
+        ),
+        (
+            {"name": "Person", "tree_root": False},
+            {"ijk": True},  # Invalid extra kwarg
+            None,
+        ),
+        (
+            ClassDefinition(name="Person", tree_root=False),
+            {"tree_root": True},
+            ClassDefinition(name="Person", tree_root=True),
+        ),
+        (
+            ClassDefinition(name="Person", tree_root=False),
+            {"ijk": True},  # Invalid extra kwarg
+            ClassDefinition(name="Person", tree_root=True),
+        ),
+    ],
+)
+def test_add_class_with_extra_kwargs(
+    cls: Union[ClassDefinition, Dict, str],
+    extra_kwargs: Dict[str, Any],
+    expected_added_class: Optional[ClassDefinition],
+):
+    """
+    Test adding a class with extra kwargs
+    """
+    # The meta slots or fields of `ClassDefinition`
+    class_meta_slots = {f.name for f in fields(ClassDefinition)}
+
+    builder = SchemaBuilder()
+
+    if not isinstance(cls, (str, dict, ClassDefinition)):
+        with pytest.raises(TypeError, match="cls must be"):
+            builder.add_class(cls, **extra_kwargs)
+    elif extra_kwargs.keys() - class_meta_slots:
+        # Handle the case of extra kwargs include a key that is not a meta slot of
+        # `ClassDefinition`
+        with pytest.raises(ValueError):
+            builder.add_class(cls, **extra_kwargs)
+    else:
+        builder.add_class(cls, **extra_kwargs)
+
+        if isinstance(cls, str):
+            class_name = cls
+        elif isinstance(cls, dict):
+            class_name = cls["name"]
+        else:
+            class_name = cls.name
+
+        added_class = builder.schema.classes[class_name]
+
+        assert added_class == expected_added_class