norm_toolkit 1.4.0__tar.gz → 1.5.0__tar.gz

{norm_toolkit-1.4.0 → norm_toolkit-1.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: norm_toolkit
-Version: 1.4.0
+Version: 1.5.0
 Summary: Toolkit to normalize text to UMLS / ontologies
 Author: Haydn Jones
 Author-email: Haydn Jones <haydnjonest@gmail.com>

{norm_toolkit-1.4.0 → norm_toolkit-1.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "norm_toolkit"
-version = "1.4.0"
+version = "1.5.0"
 description = "Toolkit to normalize text to UMLS / ontologies"
 readme = "README.md"
 authors = [{ name = "Haydn Jones", email = "haydnjonest@gmail.com" }]

{norm_toolkit-1.4.0 → norm_toolkit-1.5.0}/src/norm_toolkit/normalizer_postgres.py

@@ -917,6 +917,7 @@ ORDER BY t.concept_id, t.type_tree, t.type_id;
         concept_id: str,
         max_depth: int | None = 10,
         filter_sources: list[str] | None = None,
+        max_ids: int | None = None,
     ) -> list[str]:
         """
         Get all narrower (descendant) concept IDs using recursive traversal.
@@ -927,9 +928,11 @@ ORDER BY t.concept_id, t.type_tree, t.type_id;
             concept_id: Starting concept ID (broader term)
             max_depth: Maximum depth to traverse (1 = direct children only, None = all descendants)
             filter_sources: Only follow edges from these sources (e.g., ["SNOMEDCT_US"])
+            max_ids: Maximum number of concept IDs to return (None = no limit)
 
         Returns:
-            List of descendant concept IDs (excludes the starting concept)
+            List of descendant concept IDs ordered by depth (shallowest first),
+            excludes the starting concept
         """
         await self._ensure_initialized()
 
@@ -949,10 +952,16 @@ ORDER BY t.concept_id, t.type_tree, t.type_id;
             sources_sql = ", ".join(src_placeholders)
             source_filter = f" AND e.source IN ({sources_sql})"
 
+        # Build optional LIMIT clause
+        limit_clause = ""
+        if max_ids is not None:
+            params["max_ids"] = max_ids
+            limit_clause = "\nLIMIT :max_ids"
+
         # PostgreSQL recursive CTE with named parameters
         # Use CAST() instead of :: to avoid conflicts with SQLAlchemy named params
         # UNION (not UNION ALL) deduplicates on (concept_id, depth) during recursion
-        # DISTINCT in output needed since same concept can be reached at different depths
+        # GROUP BY with MIN(depth) gets shortest path depth for each concept
         query = f"""
 WITH RECURSIVE walk(concept_id, depth) AS (
     SELECT CAST(:concept_id AS VARCHAR), 0
@@ -964,9 +973,11 @@ WITH RECURSIVE walk(concept_id, depth) AS (
     JOIN {self._edges_table} e ON e.parent_id = w.concept_id
     WHERE (CAST(:max_depth AS INTEGER) IS NULL OR w.depth < :max_depth){source_filter}
 )
-SELECT DISTINCT concept_id
+SELECT concept_id, MIN(depth) AS min_depth
 FROM walk
 WHERE concept_id != :concept_id
+GROUP BY concept_id
+ORDER BY min_depth, concept_id{limit_clause}
 """
 
         async with self._engine.connect() as conn: