Repository: lifeomic/phc-sdk-py

Issue:
Add GenomicCopyNumber
Follow pattern of GenomicShortVariant and use this API: https://docs.us.lifeomic.com/api/#query-copy-number-data

Discussion:

Patch:
diff --git a/phc/easy/__init__.py b/phc/easy/__init__.py
index 2a2fe8c..4d90a51 100644
--- a/phc/easy/__init__.py
+++ b/phc/easy/__init__.py
@@ -20,6 +20,7 @@ from phc.easy.medication_request import MedicationRequest
 from phc.easy.medication_statement import MedicationStatement
 from phc.easy.observation import Observation
 from phc.easy.omics.genomic_short_variant import GenomicShortVariant
+from phc.easy.omics.genomic_copy_number_variant import GenomicCopyNumberVariant
 from phc.easy.omics.genomic_test import GenomicTest
 from phc.easy.organization import Organization
 from phc.easy.patients import Patient
diff --git a/phc/easy/abstract/genomic_variant.py b/phc/easy/abstract/genomic_variant.py
new file mode 100644
index 0000000..dc4d080
--- /dev/null
+++ b/phc/easy/abstract/genomic_variant.py
@@ -0,0 +1,151 @@
+import inspect
+import math
+from typing import List, Optional
+
+import pandas as pd
+from phc.easy.auth import Auth
+from phc.easy.omics.options.genomic_test import (
+    GenomicTestStatus,
+    GenomicTestType,
+)
+from phc.easy.omics.genomic_test import GenomicTest
+from phc.easy.abstract.paging_api_item import PagingApiItem
+from phc.easy.util import tqdm, split_by, rename_keys
+from phc.easy.util.batch import batch_get_frame
+
+MAX_VARIANT_SET_IDS = 100
+
+
+class GenomicVariant(PagingApiItem):
+    @staticmethod
+    def _get_genomic_tests(
+        variant_set_ids: List[str],
+        max_pages: Optional[int],
+        all_results: bool,
+        auth_args: Auth,
+        log: bool,
+        **test_args,
+    ):
+        test_args = rename_keys(test_args, {"test_status": "status"})
+        test_df = pd.DataFrame()
+        args = {}
+
+        if (
+            len(variant_set_ids) == 0
+            and max_pages is None
+            and all_results is False
+        ):
+            print("Using sample of 25 tests")
+            args = {
+                **test_args,
+                "page_size": 25,
+                "max_pages": 1,
+                "log": log,
+                "auth_args": auth_args,
+            }
+            test_df = GenomicTest.get_data_frame(**args)
+        elif len(variant_set_ids) == 0:
+            args = {
+                **test_args,
+                "page_size": 100,
+                "all_results": True,
+                "log": log,
+                "auth_args": auth_args,
+            }
+            test_df = GenomicTest.get_data_frame(**args)
+        else:
+            test_df["id"] = variant_set_ids
+
+        return test_df
+
+    @classmethod
+    def get_data_frame(
+        cls,
+        test_type: GenomicTestType,
+        # Query parameters
+        variant_set_ids: List[str] = [],
+        # Test parameters
+        patient_id: Optional[str] = None,
+        test_status: Optional[GenomicTestStatus] = GenomicTestStatus.ACTIVE,
+        # Execution parameters,
+        all_results: bool = False,
+        auth_args: Auth = Auth.shared(),
+        max_pages: Optional[int] = None,
+        page_size: Optional[int] = None,
+        log: bool = False,
+        **kw_args,
+    ):
+        """Execute a request for genomic variants
+
+        ## Parameters
+
+        Execution: `phc.easy.query.Query.execute_paging_api`
+
+        Expansion: `phc.easy.frame.Frame.expand`
+        """
+        test_params = ["patient_id", "status"]
+
+        test_args, args = split_by(
+            cls._get_current_args(inspect.currentframe(), locals()),
+            left_keys=test_params,
+        )
+
+        test_df = cls._get_genomic_tests(
+            variant_set_ids=variant_set_ids,
+            all_results=all_results,
+            test_type=test_type,
+            page_size=page_size,
+            max_pages=max_pages,
+            log=log,
+            auth_args=auth_args,
+            **test_args,
+        )
+        args["variant_set_ids"] = variant_set_ids = list(test_df.id)
+
+        if len(variant_set_ids) > MAX_VARIANT_SET_IDS and (
+            max_pages or (not all_results and page_size)
+        ):
+            print(
+                "[WARNING]: All result limit parameters are approximate when performing genomic data retrieval."
+            )
+
+        get_data_frame = super().get_data_frame
+
+        def perform_batch(ids: List[str], total_thus_far: int):
+            # Determine whether to skip this batch
+            if (
+                # Implement approximation of max_pages
+                not all_results
+                and max_pages
+                and (total_thus_far >= max_pages * (page_size or 100))
+            ) or (
+                # Use 25 or page_size for a sample (when no max_pages)
+                not all_results
+                and not max_pages
+                and total_thus_far >= (page_size or 25)
+            ):
+                return pd.DataFrame()
+
+            has_multiple_batches = len(ids) != len(variant_set_ids)
+
+            return get_data_frame(
+                **kw_args,
+                **{
+                    **args,
+                    "variant_set_ids": list(ids),
+                    "all_results": all_results
+                    # Scroll through full batches and then honor the max_pages param
+                    or (has_multiple_batches and max_pages),
+                },
+            )
+
+        variants = batch_get_frame(
+            variant_set_ids, MAX_VARIANT_SET_IDS, perform_batch
+        )
+
+        if len(variants) == 0:
+            variants["variant_set_id"] = math.nan
+
+        return variants.join(
+            test_df.set_index("id"), on="variant_set_id", rsuffix=".set"
+        )
diff --git a/phc/easy/omics/genomic_copy_number_variant.py b/phc/easy/omics/genomic_copy_number_variant.py
new file mode 100644
index 0000000..77485a1
--- /dev/null
+++ b/phc/easy/omics/genomic_copy_number_variant.py
@@ -0,0 +1,82 @@
+import inspect
+from typing import List, Optional
+import pandas as pd
+
+from phc.easy.auth import Auth
+from phc.easy.frame import Frame
+from phc.easy.omics.options.common import GenomicVariantInclude
+from phc.easy.omics.genomic_test import GenomicTestStatus, GenomicTestType
+from phc.easy.omics.options.genomic_copy_number_variant import (
+    GenomicCopyNumberVariantOptions,
+)
+from phc.easy.abstract.genomic_variant import GenomicVariant
+
+
+class GenomicCopyNumberVariant(GenomicVariant):
+    @staticmethod
+    def resource_path():
+        return "genomics/copy-numbers"
+
+    @staticmethod
+    def params_class():
+        return GenomicCopyNumberVariantOptions
+
+    @staticmethod
+    def transform_results(data_frame: pd.DataFrame, params={}, **expand_args):
+        def expand_id(id_column: pd.Series):
+            return pd.concat(
+                [
+                    id_column,
+                    id_column.str.split(":", expand=True).rename(
+                        columns={0: "variant_set_id"}
+                    )["variant_set_id"],
+                ],
+                axis=1,
+            )
+
+        args = {
+            **expand_args,
+            "custom_columns": [
+                *expand_args.get("custom_columns", []),
+                ("id", expand_id),
+            ],
+        }
+
+        return Frame.expand(data_frame, **args)
+
+    @classmethod
+    def get_data_frame(
+        cls,
+        # Query parameters
+        variant_set_ids: List[str] = [],
+        include: List[GenomicVariantInclude] = [],
+        gene: List[str] = [],
+        interpretation: List[str] = [],
+        status: List[str] = [],
+        # Test parameters
+        patient_id: Optional[str] = None,
+        test_status: Optional[GenomicTestStatus] = GenomicTestStatus.ACTIVE,
+        # Execution parameters,
+        all_results: bool = False,
+        auth_args: Auth = Auth.shared(),
+        max_pages: Optional[int] = None,
+        page_size: Optional[int] = None,
+        log: bool = False,
+        **kw_args,
+    ):
+        """Execute a request for genomic copy number variants
+
+        ## Parameters
+
+        Query: `phc.easy.omics.options.genomic_copy_number_variant.GenomicCopyNumberVariantOptions`
+
+        Execution: `phc.easy.query.Query.execute_paging_api`
+
+        Expansion: `phc.easy.frame.Frame.expand`
+        """
+
+        args = cls._get_current_args(inspect.currentframe(), locals())
+
+        return super().get_data_frame(
+            test_type=GenomicTestType.COPY_NUMBER_VARIANT, **{**kw_args, **args}
+        )
diff --git a/phc/easy/omics/genomic_short_variant.py b/phc/easy/omics/genomic_short_variant.py
index 432e263..01e7b2a 100644
--- a/phc/easy/omics/genomic_short_variant.py
+++ b/phc/easy/omics/genomic_short_variant.py
@@ -5,23 +5,20 @@ from typing import List, Optional
 import pandas as pd
 from phc.easy.auth import Auth
 from phc.easy.frame import Frame
-from phc.easy.omics.genomic_test import (
-    GenomicTest,
+from phc.easy.omics.options.genomic_test import (
     GenomicTestStatus,
     GenomicTestType,
 )
+from phc.easy.omics.options.common import GenomicVariantInclude
 from phc.easy.omics.options.genomic_short_variant import (
-    GenomicShortVariantInclude,
     GenomicShortVariantOptions,
 )
-from phc.easy.abstract.paging_api_item import PagingApiItem
-from phc.easy.util import tqdm, split_by
+from phc.easy.abstract.genomic_variant import GenomicVariant
+from phc.easy.util import split_by
 from phc.easy.util.batch import batch_get_frame
 
-MAX_VARIANT_SET_IDS = 100
 
-
-class GenomicShortVariant(PagingApiItem):
+class GenomicShortVariant(GenomicVariant):
     @staticmethod
     def resource_path():
         return "genomics/variants"
@@ -57,53 +54,12 @@ class GenomicShortVariant(PagingApiItem):
 
         return Frame.expand(data_frame, **args)
 
-    @staticmethod
-    def _get_genomic_tests(
-        variant_set_ids: List[str],
-        max_pages: Optional[int],
-        all_results: bool,
-        auth_args: Auth,
-        log: bool,
-        **test_args,
-    ):
-        test_df = pd.DataFrame()
-        args = {}
-
-        if (
-            len(variant_set_ids) == 0
-            and max_pages is None
-            and all_results is False
-        ):
-            print("Using sample of 25 tests")
-            args = {
-                **test_args,
-                "page_size": 25,
-                "max_pages": 1,
-                "log": log,
-                "auth_args": auth_args,
-            }
-        elif len(variant_set_ids) == 0:
-            args = {
-                **test_args,
-                "page_size": 100,
-                "all_results": True,
-                "log": log,
-                "auth_args": auth_args,
-            }
-        test_df = GenomicTest.get_data_frame(**args)
-
-        if len(test_df) == 0:
-            # Force id to exist as a column
-            test_df["id"] = variant_set_ids
-
-        return test_df
-
     @classmethod
     def get_data_frame(
         cls,
         # Query parameters
         variant_set_ids: List[str] = [],
-        include: List[GenomicShortVariantInclude] = ["vcf"],
+        include: List[GenomicVariantInclude] = ["vcf"],
         gene: List[str] = [],
         rsid: List[str] = [],
         chromosome: List[str] = [],
@@ -151,7 +107,7 @@ class GenomicShortVariant(PagingApiItem):
         drug_associations: Optional[bool] = None,
         # Test parameters
         patient_id: Optional[str] = None,
-        status: Optional[GenomicTestStatus] = GenomicTestStatus.ACTIVE,
+        test_status: Optional[GenomicTestStatus] = GenomicTestStatus.ACTIVE,
         # Execution parameters,
         all_results: bool = False,
         auth_args: Auth = Auth.shared(),
@@ -174,69 +130,9 @@ class GenomicShortVariant(PagingApiItem):
          - `variant_class` is translated to `class` as a parameter
          - `variant_filter` is translated to `filter` as a parameter
         """
-        TEST_PARAMS = ["patient_id", "status"]
-
-        test_args, args = split_by(
-            cls._get_current_args(inspect.currentframe(), locals()),
-            left_keys=TEST_PARAMS,
-        )
-
-        test_df = cls._get_genomic_tests(
-            variant_set_ids=variant_set_ids,
-            all_results=all_results,
-            test_type=GenomicTestType.SHORT_VARIANT,
-            page_size=page_size,
-            max_pages=max_pages,
-            log=log,
-            auth_args=auth_args,
-            **test_args,
-        )
-        variant_set_ids = list(test_df.id)
 
         args = cls._get_current_args(inspect.currentframe(), locals())
 
-        if len(variant_set_ids) > MAX_VARIANT_SET_IDS and (
-            max_pages or (not all_results and page_size)
-        ):
-            print(
-                "[WARNING]: All result limit paramters are approximate when performing genomic data retrieval."
-            )
-
-        get_data_frame = super().get_data_frame
-
-        def perform_batch(ids: List[str], total_thus_far: int):
-            # Determine whether to skip this batch
-            if (
-                # Implement approximation of max_pages
-                not all_results
-                and max_pages
-                and (total_thus_far >= max_pages * (page_size or 100))
-            ) or (
-                # Use 25 or page_size for a sample (when no max_pages)
-                not all_results
-                and not max_pages
-                and total_thus_far >= (page_size or 25)
-            ):
-                return pd.DataFrame()
-
-            has_multiple_batches = len(ids) != len(variant_set_ids)
-
-            return get_data_frame(
-                **kw_args,
-                **{
-                    **args,
-                    "variant_set_ids": list(ids),
-                    "all_results": all_results
-                    # Scroll through full batches and then honor the max_pages param
-                    or (has_multiple_batches and max_pages),
-                },
-            )
-
-        variants = batch_get_frame(
-            variant_set_ids, MAX_VARIANT_SET_IDS, perform_batch
+        return super().get_data_frame(
+            test_type=GenomicTestType.SHORT_VARIANT, **{**kw_args, **args}
         )
-
-        if len(variants) == 0:
-            variants["variant_set_id"] = math.nan
-
-        return variants.join(test_df.set_index("id"), on="variant_set_id")
diff --git a/phc/easy/omics/options/common.py b/phc/easy/omics/options/common.py
new file mode 100644
index 0000000..0718ffa
--- /dev/null
+++ b/phc/easy/omics/options/common.py
@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class GenomicVariantInclude(str, Enum):
+    ENSEMBL = "ensembl"
+    VCF = "vcf"
+    DRUG_ASSOCIATIONS = "drugAssociations"
diff --git a/phc/easy/omics/options/genomic_copy_number_variant.py b/phc/easy/omics/options/genomic_copy_number_variant.py
new file mode 100644
index 0000000..a7e015c
--- /dev/null
+++ b/phc/easy/omics/options/genomic_copy_number_variant.py
@@ -0,0 +1,32 @@
+from enum import Enum
+from typing import List, Optional
+
+from phc.easy.omics.options.common import GenomicVariantInclude
+from phc.easy.abstract.paging_api_item import PagingApiOptions
+from pydantic import Field
+
+MAPPINGS = {
+    "variant_set_ids": "copyNumberSetIds",
+    "drug_associations": "drugAssociations",
+}
+
+
+class GenomicCopyNumberVariantOptions(PagingApiOptions):
+    """Options to pass to `/v1/genomics/copy-numbers`
+
+    See https://docs.us.lifeomic.com/api/#query-copy-number-data
+    """
+
+    variant_set_ids: List[str] = Field(..., min_items=1)
+    include: List[GenomicVariantInclude] = []
+    gene: List[str] = []
+    interpretation: List[str] = []
+    status: List[str] = []
+    drug_associations: Optional[bool]
+
+    @staticmethod
+    def transform(key, value):
+        if key not in ["drug_associations"]:
+            value = ",".join(value)
+
+        return (MAPPINGS.get(key, key), value)
diff --git a/phc/easy/omics/options/genomic_short_variant.py b/phc/easy/omics/options/genomic_short_variant.py
index 6962d54..6ef7aa4 100644
--- a/phc/easy/omics/options/genomic_short_variant.py
+++ b/phc/easy/omics/options/genomic_short_variant.py
@@ -1,6 +1,7 @@
 from enum import Enum
 from typing import List, Optional
 
+from phc.easy.omics.options.common import GenomicVariantInclude
 from phc.easy.abstract.paging_api_item import PagingApiOptions
 from pydantic import Field
 
@@ -44,12 +45,6 @@ MAPPINGS = {
 }
 
 
-class GenomicShortVariantInclude(str, Enum):
-    ENSEMBL = "ensembl"
-    VCF = "vcf"
-    DRUG_ASSOCIATIONS = "drugAssociations"
-
-
 class GenomicShortVariantOptions(PagingApiOptions):
     """Options to pass to `/v1/genomics/variants`
 
@@ -57,7 +52,7 @@ class GenomicShortVariantOptions(PagingApiOptions):
     """
 
     variant_set_ids: List[str] = Field(..., min_items=1)
-    include: List[GenomicShortVariantInclude] = ["vcf"]
+    include: List[GenomicVariantInclude] = ["vcf"]
     gene: List[str] = []
     rsid: List[str] = []
     chromosome: List[str] = []
diff --git a/phc/easy/util/__init__.py b/phc/easy/util/__init__.py
index daa6ebb..457a935 100644
--- a/phc/easy/util/__init__.py
+++ b/phc/easy/util/__init__.py
@@ -15,6 +15,11 @@ else:
     _has_tqdm = True
 
 
+def rename_keys(dictionary: dict, mapping: dict):
+    "Rename keys in a dictionary"
+    return {mapping.get(k, k): v for k, v in dictionary.items()}
+
+
 def join_underscore(values):
     return "_".join(
         [


Test patch:
diff --git a/phc/easy/omics/genomic_test.py b/phc/easy/omics/genomic_test.py
index af33ccc..31ec79c 100644
--- a/phc/easy/omics/genomic_test.py
+++ b/phc/easy/omics/genomic_test.py
@@ -5,23 +5,13 @@ from typing import Optional
 import pandas as pd
 from phc.easy.auth import Auth
 from phc.easy.frame import Frame
+from phc.easy.omics.options.genomic_test import (
+    GenomicTestStatus,
+    GenomicTestType,
+)
 from phc.easy.abstract.paging_api_item import PagingApiItem, PagingApiOptions
 
 
-class GenomicTestType(str, Enum):
-    SHORT_VARIANT = "shortVariant"
-    EXPRESSION = "expression"
-    STRUCTURAL_VARIANT = "structuralVariant"
-    COPY_NUMBER_VARIANT = "copyNumberVariant"
-    READ = "read"
-
-
-class GenomicTestStatus(str, Enum):
-    ACTIVE = "ACTIVE"
-    INDEXING = "INDEXING"
-    FAILED = "FAILED"
-
-
 class GenomicTestOptions(PagingApiOptions):
     """Options to pass to `/v1/genomics/projects/:project_id/tests`
 
diff --git a/phc/easy/omics/options/genomic_test.py b/phc/easy/omics/options/genomic_test.py
new file mode 100644
index 0000000..31e1693
--- /dev/null
+++ b/phc/easy/omics/options/genomic_test.py
@@ -0,0 +1,15 @@
+from enum import Enum
+
+
+class GenomicTestType(str, Enum):
+    SHORT_VARIANT = "shortVariant"
+    EXPRESSION = "expression"
+    STRUCTURAL_VARIANT = "structuralVariant"
+    COPY_NUMBER_VARIANT = "copyNumberVariant"
+    READ = "read"
+
+
+class GenomicTestStatus(str, Enum):
+    ACTIVE = "ACTIVE"
+    INDEXING = "INDEXING"
+    FAILED = "FAILED"
diff --git a/tests/test_easy_util.py b/tests/test_easy_util.py
index 34f4859..306928e 100644
--- a/tests/test_easy_util.py
+++ b/tests/test_easy_util.py
@@ -1,3 +1,5 @@
+from nose.tools import assert_equals
+
 from phc.easy.util import (
     add_prefixes,
     concat_dicts,
@@ -7,6 +9,7 @@ from phc.easy.util import (
     prefix_dict_keys,
     without_keys,
     split_by,
+    rename_keys,
 )
 
 
@@ -110,6 +113,10 @@ def test_split_by():
     assert list(set_b.keys()) == ["b", "d"]
 
 
+def test_rename_keys():
+    assert_equals(rename_keys({"a": 1, "b": 2}, {"a": "aa"}), {"aa": 1, "b": 2})
+
+
 # defaultprop
 
 
diff --git a/tests/test_genomic_copy_number_variant.py b/tests/test_genomic_copy_number_variant.py
new file mode 100644
index 0000000..00491d6
--- /dev/null
+++ b/tests/test_genomic_copy_number_variant.py
@@ -0,0 +1,25 @@
+import pandas as pd
+from nose.tools import assert_equals
+from phc.easy.omics.genomic_copy_number_variant import GenomicCopyNumberVariant
+
+
+def test_parse_id():
+    raw_df = pd.DataFrame(
+        # NOTE: Sample is taken and adapted from BRCA data set
+        [
+            {"id": "6b0591ce-7b3b-4b04-85bc-d17e463ca869:A235y+Jw+v="},
+            {"id": "f0e381b6-a9b3-4411-af56-7f7f5ce3ce6b:XjQLzpOuLm="},
+            {"id": "6b0591ce-7b3b-4b04-85bc-d17e463ca869:BsZ3G0NtUz="},
+        ]
+    )
+    frame = GenomicCopyNumberVariant.transform_results(raw_df)
+
+    assert_equals(list(frame.columns), ["id", "variant_set_id"])
+
+    assert_equals(
+        list(frame.variant_set_id.unique()),
+        [
+            "6b0591ce-7b3b-4b04-85bc-d17e463ca869",
+            "f0e381b6-a9b3-4411-af56-7f7f5ce3ce6b",
+        ],
+    )
diff --git a/tests/test_genomic_short_variant.py b/tests/test_genomic_short_variant.py
index 69ec6d0..eb614ef 100644
--- a/tests/test_genomic_short_variant.py
+++ b/tests/test_genomic_short_variant.py
@@ -2,6 +2,7 @@ import pandas as pd
 from nose.tools import assert_equals
 from unittest import mock
 from unittest.mock import ANY
+from phc.easy.auth import Auth
 from phc.easy.omics.genomic_short_variant import GenomicShortVariant
 from uuid import uuid4
 
@@ -52,3 +53,23 @@ def test_batches_of_variant_set_ids(execute_paging_api, test_get_data_frame):
     assert_equals(len(get_variant_set_ids(0)), 100)
     assert_equals(len(get_variant_set_ids(1)), 100)
     assert_equals(len(get_variant_set_ids(2)), 50)
+
+    # GenomicTest should not be retrieved
+    assert_equals(test_get_data_frame.call_count, 0)
+
+
+@mock.patch("phc.easy.query.Query.execute_paging_api")
+def test_passing_options_through_to_paging_api(execute_paging_api):
+    execute_paging_api.return_value = pd.DataFrame()
+
+    auth = Auth()
+
+    GenomicShortVariant.get_data_frame(
+        [str(uuid4())], raw=True, log=True, auth_args=auth
+    )
+
+    kwargs = execute_paging_api.call_args[1]
+
+    assert_equals(kwargs.get("auth_args"), auth)
+    assert_equals(kwargs.get("log"), True)
+    assert_equals(kwargs.get("raw"), True)
diff --git a/tests/test_genomic_variant.py b/tests/test_genomic_variant.py
new file mode 100644
index 0000000..73a2e75
--- /dev/null
+++ b/tests/test_genomic_variant.py
@@ -0,0 +1,32 @@
+from uuid import uuid4
+from unittest import mock
+from nose.tools import assert_equals
+from phc.easy.auth import Auth
+from phc.easy.abstract.genomic_variant import GenomicVariant
+
+
+@mock.patch("phc.easy.omics.genomic_test.GenomicTest.get_data_frame")
+def test_getting_genomic_tests(get_data_frame):
+    GenomicVariant._get_genomic_tests(
+        [], max_pages=None, all_results=False, auth_args=Auth(), log=False
+    )
+
+    get_data_frame.assert_called_once()
+
+
+@mock.patch("phc.easy.omics.genomic_test.GenomicTest.get_data_frame")
+def test_skipping_genomic_tests_if_variant_set_ids(get_data_frame):
+    variant_set_ids = [str(uuid4())]
+
+    test_df = GenomicVariant._get_genomic_tests(
+        variant_set_ids,
+        max_pages=None,
+        all_results=False,
+        auth_args=Auth(),
+        log=False,
+    )
+
+    assert_equals(get_data_frame.call_count, 0)
+
+    assert_equals(len(test_df.columns), 1)
+    assert_equals(list(test_df.id), variant_set_ids)