From b8d2e4cf8309fa2d92cca4a9e95680664bb7c4e6 Mon Sep 17 00:00:00 2001
From: Charles Tapley Hoyt <cthoyt@gmail.com>
Date: Sat, 23 Nov 2024 13:55:45 +0100
Subject: [PATCH] Update struct.py

---
 src/pyobo/api/hierarchy.py | 24 +++++++++++++++---------
 src/pyobo/api/relations.py |  7 +++++--
 src/pyobo/sources/rhea.py  |  3 +--
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/src/pyobo/api/hierarchy.py b/src/pyobo/api/hierarchy.py
index 92de9398..816ab626 100644
--- a/src/pyobo/api/hierarchy.py
+++ b/src/pyobo/api/hierarchy.py
@@ -10,8 +10,9 @@
 from .properties import get_filtered_properties_mapping
 from .relations import get_filtered_relations_df
 from ..identifier_utils import wrap_norm_prefix
-from ..struct import TypeDef, has_member, is_a, part_of
+from ..struct import has_member, is_a, part_of
 from ..struct.reference import Reference
+from ..struct.struct import ReferenceHint, _ensure_ref
 
 __all__ = [
     "get_ancestors",
@@ -32,8 +33,8 @@ def get_hierarchy(
     *,
     include_part_of: bool = True,
     include_has_member: bool = False,
-    extra_relations: Iterable[TypeDef] | None = None,
-    properties: Iterable[str] | None = None,
+    extra_relations: Iterable[ReferenceHint] | None = None,
+    properties: Iterable[ReferenceHint] | None = None,
     use_tqdm: bool = False,
     force: bool = False,
     force_process: bool = False,
@@ -58,12 +59,19 @@ def get_hierarchy(
 
     This function thinly wraps :func:`_get_hierarchy_helper` to make it easier to work with the lru_cache mechanism.
     """
+    extra_relations_ = tuple(
+        sorted(_ensure_ref(r, ontology_prefix=prefix) for r in extra_relations or [])
+    )
+    properties_ = tuple(
+        sorted(_ensure_ref(prop, ontology_prefix=prefix) for prop in properties or [])
+    )
+
     return _get_hierarchy_helper(
         prefix=prefix,
         include_part_of=include_part_of,
         include_has_member=include_has_member,
-        extra_relations=tuple(sorted(extra_relations or [])),
-        properties=tuple(sorted(properties or [])),
+        extra_relations=extra_relations_,
+        properties=properties_,
         use_tqdm=use_tqdm,
         force=force,
         force_process=force_process,
@@ -77,8 +85,8 @@ def get_hierarchy(
 def _get_hierarchy_helper(
     prefix: str,
     *,
-    extra_relations: tuple[TypeDef, ...],
-    properties: tuple[str, ...],
+    extra_relations: tuple[Reference, ...],
+    properties: tuple[Reference, ...],
     include_part_of: bool,
     include_has_member: bool,
     use_tqdm: bool,
@@ -140,8 +148,6 @@ def _get_hierarchy_helper(
             rv.add_edge(f"{source_ns}:{source_id}", f"{prefix}:{target_id}", relation="part_of")
 
     for relation in extra_relations:
-        if not isinstance(relation, TypeDef | Reference):
-            raise TypeError
         relation_df = get_filtered_relations_df(
             prefix=prefix,
             relation=relation,
diff --git a/src/pyobo/api/relations.py b/src/pyobo/api/relations.py
index e1bbf3ef..ea7ce11c 100644
--- a/src/pyobo/api/relations.py
+++ b/src/pyobo/api/relations.py
@@ -22,7 +22,6 @@
 )
 from ..getters import get_ontology
 from ..identifier_utils import wrap_norm_prefix
-from ..struct import TypeDef
 from ..struct.reference import Reference
 from ..struct.struct import ReferenceHint, _ensure_ref
 from ..utils.cache import cached_df
@@ -102,7 +101,11 @@ def _df_getter() -> pd.DataFrame:
 
 @wrap_norm_prefix
 def get_id_multirelations_mapping(
-    prefix: str, typedef: TypeDef, *, use_tqdm: bool = False, **kwargs: Unpack[GetOntologyKwargs]
+    prefix: str,
+    typedef: ReferenceHint,
+    *,
+    use_tqdm: bool = False,
+    **kwargs: Unpack[GetOntologyKwargs],
 ) -> Mapping[str, list[Reference]]:
     """Get the OBO file and output a synonym dictionary."""
     kwargs["version"] = get_version_from_kwargs(prefix, kwargs)
diff --git a/src/pyobo/sources/rhea.py b/src/pyobo/sources/rhea.py
index aaecf9f0..edbf03ac 100644
--- a/src/pyobo/sources/rhea.py
+++ b/src/pyobo/sources/rhea.py
@@ -9,7 +9,6 @@
 from pyobo.api.utils import get_version
 from pyobo.struct import Obo, Reference, Term
 from pyobo.struct.typedef import (
-    TypeDef,
     enabled_by,
     has_bidirectional_reaction,
     has_input,
@@ -208,7 +207,7 @@ def iter_terms(version: str, force: bool = False) -> Iterable[Term]:
                 )
                 continue
             target_reference = Reference(prefix=xref_prefix, identifier=xref_id)
-            if isinstance(relation, TypeDef):
+            if relation is not None:
                 terms[directional_rhea_id].append_relationship(relation, target_reference)
             else:
                 terms[directional_rhea_id].append_xref(target_reference)