Extending mapping method with source_codes option

Author: alex-omophub

src/omophub/resources/mappings.py

@@ -47,29 +47,50 @@ def get(
 
     def map(
         self,
-        source_concepts: list[int],
         target_vocabulary: str,
         *,
+        source_concepts: list[int] | None = None,
+        source_codes: list[dict[str, str]] | None = None,
         mapping_type: str | None = None,
         include_invalid: bool = False,
         vocab_release: str | None = None,
     ) -> dict[str, Any]:
         """Map concepts to a target vocabulary.
 
         Args:
-            source_concepts: List of OMOP concept IDs to map
-            target_vocabulary: Target vocabulary ID (e.g., "ICD10CM", "SNOMED")
-            mapping_type: Mapping type (direct, equivalent, broader, narrower)
+            target_vocabulary: Target vocabulary ID (e.g., "ICD10CM", "SNOMED", "RxNorm")
+            source_concepts: List of OMOP concept IDs to map. Use this OR source_codes,
+                not both.
+            source_codes: List of vocabulary/code pairs to map, e.g.,
+                [{"vocabulary_id": "SNOMED", "concept_code": "387517004"}].
+                Use this OR source_concepts, not both.
+            mapping_type: Mapping type filter (direct, equivalent, broader, narrower)
             include_invalid: Include invalid mappings
             vocab_release: Specific vocabulary release version (e.g., "2025.1")
 
         Returns:
             Mapping results with summary
+
+        Raises:
+            ValueError: If neither or both source_concepts and source_codes are provided
         """
+        # Validate: exactly one of source_concepts or source_codes required
+        has_concepts = source_concepts is not None and len(source_concepts) > 0
+        has_codes = source_codes is not None and len(source_codes) > 0
+
+        if not has_concepts and not has_codes:
+            raise ValueError("Either source_concepts or source_codes is required")
+        if has_concepts and has_codes:
+            raise ValueError("Cannot use both source_concepts and source_codes")
+
         body: dict[str, Any] = {
-            "source_concepts": source_concepts,
             "target_vocabulary": target_vocabulary,
         }
+
+        if source_concepts:
+            body["source_concepts"] = source_concepts
+        if source_codes:
+            body["source_codes"] = source_codes
         if mapping_type:
             body["mapping_type"] = mapping_type
         if include_invalid:
@@ -123,29 +144,50 @@ async def get(
 
     async def map(
         self,
-        source_concepts: list[int],
         target_vocabulary: str,
         *,
+        source_concepts: list[int] | None = None,
+        source_codes: list[dict[str, str]] | None = None,
         mapping_type: str | None = None,
         include_invalid: bool = False,
         vocab_release: str | None = None,
     ) -> dict[str, Any]:
         """Map concepts to a target vocabulary.
 
         Args:
-            source_concepts: List of OMOP concept IDs to map
-            target_vocabulary: Target vocabulary ID (e.g., "ICD10CM", "SNOMED")
-            mapping_type: Mapping type (direct, equivalent, broader, narrower)
+            target_vocabulary: Target vocabulary ID (e.g., "ICD10CM", "SNOMED", "RxNorm")
+            source_concepts: List of OMOP concept IDs to map. Use this OR source_codes,
+                not both.
+            source_codes: List of vocabulary/code pairs to map, e.g.,
+                [{"vocabulary_id": "SNOMED", "concept_code": "387517004"}].
+                Use this OR source_concepts, not both.
+            mapping_type: Mapping type filter (direct, equivalent, broader, narrower)
             include_invalid: Include invalid mappings
             vocab_release: Specific vocabulary release version (e.g., "2025.1")
 
         Returns:
             Mapping results with summary
+
+        Raises:
+            ValueError: If neither or both source_concepts and source_codes are provided
         """
+        # Validate: exactly one of source_concepts or source_codes required
+        has_concepts = source_concepts is not None and len(source_concepts) > 0
+        has_codes = source_codes is not None and len(source_codes) > 0
+
+        if not has_concepts and not has_codes:
+            raise ValueError("Either source_concepts or source_codes is required")
+        if has_concepts and has_codes:
+            raise ValueError("Cannot use both source_concepts and source_codes")
+
         body: dict[str, Any] = {
-            "source_concepts": source_concepts,
             "target_vocabulary": target_vocabulary,
         }
+
+        if source_concepts:
+            body["source_concepts"] = source_concepts
+        if source_codes:
+            body["source_codes"] = source_codes
         if mapping_type:
             body["mapping_type"] = mapping_type
         if include_invalid:

tests/unit/resources/test_mappings.py

@@ -79,8 +79,8 @@ def test_map_concepts(self, sync_client: OMOPHub, base_url: str) -> None:
         )
 
         result = sync_client.mappings.map(
-            source_concepts=[{"concept_id": 201826}],
             target_vocabulary="ICD10CM",
+            source_concepts=[201826],
         )
 
         assert "mappings" in result
@@ -97,18 +97,72 @@ def test_map_concepts_with_options(
         )
 
         sync_client.mappings.map(
-            source_concepts=[
-                {"concept_id": 201826},
-                {"vocabulary_id": "SNOMED", "concept_code": "44054006"},
-            ],
             target_vocabulary="ICD10CM",
+            source_concepts=[201826, 4329847],
             mapping_type="equivalent",
             include_invalid=True,
         )
 
         # Verify POST body
         assert route.calls[0].request.content
 
+    @respx.mock
+    def test_map_concepts_with_source_codes(
+        self, sync_client: OMOPHub, base_url: str
+    ) -> None:
+        """Test mapping concepts using source_codes parameter."""
+        import json
+
+        map_response = {
+            "success": True,
+            "data": {
+                "mappings": [
+                    {
+                        "source_concept_id": 4306040,
+                        "source_concept_name": "Acetaminophen",
+                        "target_concept_id": 1125315,
+                        "target_vocabulary_id": "RxNorm",
+                    }
+                ],
+            },
+        }
+        route = respx.post(f"{base_url}/concepts/map").mock(
+            return_value=Response(200, json=map_response)
+        )
+
+        result = sync_client.mappings.map(
+            target_vocabulary="RxNorm",
+            source_codes=[
+                {"vocabulary_id": "SNOMED", "concept_code": "387517004"},
+                {"vocabulary_id": "SNOMED", "concept_code": "108774000"},
+            ],
+        )
+
+        assert "mappings" in result
+        # Verify request body contains source_codes
+        body = json.loads(route.calls[0].request.content)
+        assert "source_codes" in body
+        assert len(body["source_codes"]) == 2
+        assert body["source_codes"][0]["vocabulary_id"] == "SNOMED"
+
+    def test_map_concepts_requires_source(self, sync_client: OMOPHub) -> None:
+        """Test that map() raises error when neither source_concepts nor source_codes provided."""
+        with pytest.raises(
+            ValueError, match="Either source_concepts or source_codes is required"
+        ):
+            sync_client.mappings.map(target_vocabulary="ICD10CM")
+
+    def test_map_concepts_rejects_both_sources(self, sync_client: OMOPHub) -> None:
+        """Test that map() raises error when both source_concepts and source_codes provided."""
+        with pytest.raises(
+            ValueError, match="Cannot use both source_concepts and source_codes"
+        ):
+            sync_client.mappings.map(
+                target_vocabulary="ICD10CM",
+                source_concepts=[201826],
+                source_codes=[{"vocabulary_id": "SNOMED", "concept_code": "44054006"}],
+            )
+
 
 class TestAsyncMappingsResource:
     """Tests for the asynchronous AsyncMappings resource."""
@@ -157,8 +211,8 @@ async def test_async_map_concepts(
         )
 
         result = await async_client.mappings.map(
-            source_concepts=[{"concept_id": 201826}],
             target_vocabulary="ICD10CM",
+            source_concepts=[201826],
         )
 
         assert "mappings" in result
@@ -174,10 +228,57 @@ async def test_async_map_concepts_with_options(
         )
 
         await async_client.mappings.map(
-            source_concepts=[{"concept_id": 201826}],
             target_vocabulary="ICD10CM",
+            source_concepts=[201826],
             mapping_type="direct",
             include_invalid=True,
         )
 
         assert route.calls[0].request.content
+
+    @pytest.mark.asyncio
+    @respx.mock
+    async def test_async_map_concepts_with_source_codes(
+        self, async_client: omophub.AsyncOMOPHub, base_url: str
+    ) -> None:
+        """Test async mapping concepts using source_codes."""
+        import json
+
+        route = respx.post(f"{base_url}/concepts/map").mock(
+            return_value=Response(200, json={"success": True, "data": {"mappings": []}})
+        )
+
+        result = await async_client.mappings.map(
+            target_vocabulary="RxNorm",
+            source_codes=[
+                {"vocabulary_id": "SNOMED", "concept_code": "387517004"},
+            ],
+        )
+
+        assert "mappings" in result
+        body = json.loads(route.calls[0].request.content)
+        assert "source_codes" in body
+
+    @pytest.mark.asyncio
+    async def test_async_map_requires_source(
+        self, async_client: omophub.AsyncOMOPHub
+    ) -> None:
+        """Test async map() raises error without sources."""
+        with pytest.raises(
+            ValueError, match="Either source_concepts or source_codes is required"
+        ):
+            await async_client.mappings.map(target_vocabulary="ICD10CM")
+
+    @pytest.mark.asyncio
+    async def test_async_map_rejects_both_sources(
+        self, async_client: omophub.AsyncOMOPHub
+    ) -> None:
+        """Test async map() raises error with both sources."""
+        with pytest.raises(
+            ValueError, match="Cannot use both source_concepts and source_codes"
+        ):
+            await async_client.mappings.map(
+                target_vocabulary="ICD10CM",
+                source_concepts=[201826],
+                source_codes=[{"vocabulary_id": "SNOMED", "concept_code": "44054006"}],
+            )