diff --git a/brickschema/inference.py b/brickschema/inference.py index 8d53bae..ad7a557 100644 --- a/brickschema/inference.py +++ b/brickschema/inference.py @@ -730,7 +730,7 @@ def infer_model(self, model): entities = model["rows"] # index the entities by their ID field - entities = {e["id"].replace('"', ""): {"tags": e} for e in entities} + entities = {_extract_id(e).replace('"', ""): {"tags": e} for e in entities} # TODO: add e['dis'] for a descriptive label? brickgraph = Graph(load_brick=False) @@ -744,7 +744,7 @@ def infer_model(self, model): # translate tags entity_tagset = list(self._translate_tags(marker_tags)) - equip_ref = entity["tags"].get("equipRef") + equip_ref = _extract_ref(entity["tags"], "equipRef") # infer tags for single entity triples, _ = self.infer_entity( entity_tagset, identifier=entity_id, equip_ref=equip_ref @@ -760,7 +760,8 @@ def infer_model(self, model): if "equipRef" not in relships: continue reffed_equip = ( - relships["equipRef"].replace(" ", "_").replace('"', "") + "_equip" + _extract_ref(relships, "equipRef").replace(" ", "_").replace('"', "") + + "_equip" ) if self._BLDG[point_entity_id] in brickgraph.nodes: triple = ( @@ -803,3 +804,23 @@ def _to_tag_case(x): x (str): transformed string """ return x[0].upper() + x[1:] + + +def _extract_id(entity): + """ + Extracts the 'id' field from Haystack entity + """ + return _extract_ref(entity, "id") + + +def _extract_ref(entity, ref): + """ + Extracts the indicated 'ref' field from Haystack entity + """ + if ref not in entity: + return None + if isinstance(entity[ref], str): + return entity[ref].replace('"', "") # strip quotes + elif isinstance(entity[ref], dict): + return entity[ref]["val"].replace('"', "") + return None