refactor: Move template inference logic from Library to ShapeCollection

NREL · Dec 17, 2024 · f525260 · f525260
1 parent 1ae16e0
commit f525260
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 41 deletions.
diff --git a/buildingmotif/dataclasses/library.py b/buildingmotif/dataclasses/library.py
@@ -285,54 +285,18 @@ def _load_from_ontology(
 
         lib = cls.create(ontology_name, overwrite=overwrite)
 
-        if infer_templates:
-            # infer shapes from any class/nodeshape candidates in the graph
-            lib._infer_templates_from_graph(ontology)
-
         # load the ontology graph as a shape_collection
         shape_col_id = lib.get_shape_collection().id
         assert shape_col_id is not None  # should always pass
         shape_col = ShapeCollection.load(shape_col_id)
         shape_col.add_graph(ontology)
 
-        return lib
-
-    def _infer_templates_from_graph(self, graph: rdflib.Graph):
-        """Infer templates from a graph (by interpreting shapes) and add them to this library.
+        if infer_templates:
+            # infer shapes from any class/nodeshape candidates in the graph
+            shape_col.infer_templates(lib)
 
-        :param graph: graph to infer templates from
-        :type graph: rdflib.Graph
-        """
-        # add all imports to the same graph so we can resolve everything
-        imports_closure = copy_graph(graph)
-        # import dependencies into 'graph'
-        # get all imports from the graph
-        for dependency in graph.objects(predicate=rdflib.OWL.imports):
-            # attempt to load from BuildingMOTIF
-            try:
-                lib = Library.load(name=str(dependency))
-                imports_closure += lib.get_shape_collection().graph
-            except Exception as e:  # TODO: replace with a more specific exception
-                logging.warning(
-                    f"An ontology could not resolve a dependency on {dependency} ({e}). Check this is loaded into BuildingMOTIF"
-                )
-                continue
-        class_candidates = set(graph.subjects(rdflib.RDF.type, rdflib.OWL.Class))
-        shape_candidates = set(graph.subjects(rdflib.RDF.type, rdflib.SH.NodeShape))
-        candidates = class_candidates.intersection(shape_candidates)
-        template_id_lookup: Dict[str, int] = {}
-        dependency_cache: Dict[int, List[Dict[Any, Any]]] = {}
-        for candidate in candidates:
-            assert isinstance(candidate, rdflib.URIRef)
-            # TODO: mincount 0 (or unspecified) should be optional args on the generated template
-            partial_body, deps = get_template_parts_from_shape(
-                candidate, imports_closure
-            )
-            templ = self.create_template(str(candidate), partial_body)
-            dependency_cache[templ.id] = deps
-            template_id_lookup[str(candidate)] = templ.id
+        return lib
 
-        self._resolve_template_dependencies(template_id_lookup, dependency_cache)
 
     def _load_shapes_from_directory(
         self,
@@ -367,7 +331,7 @@ def _load_shapes_from_directory(
             )
         # infer shapes from any class/nodeshape candidates in the graph
         if infer_templates:
-            self._infer_templates_from_graph(shape_col.graph)
+            shape_col.infer_templates(self)
 
     @classmethod
     def _load_from_directory(

diff --git a/buildingmotif/dataclasses/shape_collection.py b/buildingmotif/dataclasses/shape_collection.py
@@ -180,6 +180,39 @@ def _get_included_domains(cls, domain: URIRef) -> List[URIRef]:
 
         return results
 
+    def infer_templates(self, library: "Library") -> None:
+        """Infer templates from the graph in this ShapeCollection and add them to the given library.
+
+        :param library: The library to add inferred templates to
+        :type library: Library
+        """
+        imports_closure = copy_graph(self.graph)
+        for dependency in self.graph.objects(predicate=rdflib.OWL.imports):
+            try:
+                lib = Library.load(name=str(dependency))
+                imports_closure += lib.get_shape_collection().graph
+            except Exception as e:
+                logging.warning(
+                    f"An ontology could not resolve a dependency on {dependency} ({e}). Check this is loaded into BuildingMOTIF"
+                )
+                continue
+
+        class_candidates = set(self.graph.subjects(rdflib.RDF.type, rdflib.OWL.Class))
+        shape_candidates = set(self.graph.subjects(rdflib.RDF.type, rdflib.SH.NodeShape))
+        candidates = class_candidates.intersection(shape_candidates)
+
+        template_id_lookup: Dict[str, int] = {}
+        dependency_cache: Dict[int, List[Dict[Any, Any]]] = {}
+
+        for candidate in candidates:
+            assert isinstance(candidate, rdflib.URIRef)
+            partial_body, deps = get_template_parts_from_shape(candidate, imports_closure)
+            templ = library.create_template(str(candidate), partial_body)
+            dependency_cache[templ.id] = deps
+            template_id_lookup[str(candidate)] = templ.id
+
+        library._resolve_template_dependencies(template_id_lookup, dependency_cache)
+
     def get_shapes_of_definition_type(
         self, definition_type: URIRef, include_labels=False
     ) -> Union[List[URIRef], List[Tuple[URIRef, str]]]: