Repository: lifeomic/phc-sdk-py

Issue:
Add max_pages for PatientItem
Especially when dealing with a tremendous amount of data, we'd like to page until we reach a certain count.

Comments:

Patch:
diff --git a/phc/easy/item.py b/phc/easy/item.py
index b764b79..77c4b90 100644
--- a/phc/easy/item.py
+++ b/phc/easy/item.py
@@ -45,6 +45,7 @@ class Item:
         cls,
         all_results: bool = False,
         raw: bool = False,
+        max_pages: Union[int, None] = None,
         query_overrides: dict = {},
         auth_args=Auth.shared(),
         ignore_cache: bool = False,
@@ -62,6 +63,9 @@ class Item:
             If raw, then values will not be expanded (useful for manual
             inspection if something goes wrong)
 
+        max_pages : int
+            The number of pages to retrieve (useful if working with tons of records)
+
         query_overrides : dict = {}
             Override any part of the elasticsearch FHIR query
 
@@ -105,6 +109,7 @@ class Item:
             query_overrides,
             auth_args,
             ignore_cache,
+            max_pages=max_pages,
             log=log,
         )
 
diff --git a/phc/easy/patient_item.py b/phc/easy/patient_item.py
index ccf1f46..96e6a15 100644
--- a/phc/easy/patient_item.py
+++ b/phc/easy/patient_item.py
@@ -28,6 +28,7 @@ class PatientItem(Item):
         raw: bool = False,
         patient_id: Union[None, str] = None,
         patient_ids: List[str] = [],
+        max_pages: Union[int, None] = None,
         query_overrides: dict = {},
         auth_args=Auth.shared(),
         ignore_cache: bool = False,
@@ -51,6 +52,9 @@ class PatientItem(Item):
         patient_ids : List[str]
             Find records for given patient_ids
 
+        max_pages : int
+            The number of pages to retrieve (useful if working with tons of records)
+
         query_overrides : dict = {}
             Override any part of the elasticsearch FHIR query
 
@@ -96,6 +100,7 @@ class PatientItem(Item):
             ignore_cache,
             patient_id=patient_id,
             patient_ids=patient_ids,
+            max_pages=max_pages,
             patient_key=cls.patient_key(),
             log=log,
             patient_id_prefixes=cls.patient_id_prefixes(),
diff --git a/phc/easy/patients/__init__.py b/phc/easy/patients/__init__.py
index c6ecfb0..fe3c26d 100644
--- a/phc/easy/patients/__init__.py
+++ b/phc/easy/patients/__init__.py
@@ -1,24 +1,23 @@
 import pandas as pd
 
-from phc.easy.auth import Auth
 from phc.easy.frame import Frame
+from phc.easy.item import Item
 from phc.easy.patients.address import expand_address_column
 from phc.easy.patients.name import expand_name_column
-from phc.easy.query import Query
 
 
-class Patient:
+class Patient(Item):
     @staticmethod
-    def get_count(query_overrides: dict = {}, auth_args=Auth.shared()):
-        return Query.find_count_of_dsl_query(
-            {
-                "type": "select",
-                "columns": "*",
-                "from": [{"table": "patient"}],
-                **query_overrides,
-            },
-            auth_args=auth_args,
-        )
+    def table_name():
+        return "patient"
+
+    @staticmethod
+    def code_keys():
+        return [
+            "extension.valueCodeableConcept.coding",
+            "identifier.type.coding",
+            "meta.tag",
+        ]
 
     @staticmethod
     def transform_results(data_frame: pd.DataFrame, **expand_args):
@@ -32,65 +31,3 @@ class Patient:
         }
 
         return Frame.expand(data_frame, **args)
-
-    @staticmethod
-    def get_data_frame(
-        all_results: bool = False,
-        raw: bool = False,
-        query_overrides: dict = {},
-        auth_args: Auth = Auth.shared(),
-        ignore_cache: bool = False,
-        expand_args: dict = {},
-    ):
-        """Retrieve patients as a data frame with unwrapped FHIR columns
-
-        Attributes
-        ----------
-        all_results : bool = False
-            Override limit to retrieve all patients
-
-        raw : bool = False
-            If raw, then values will not be expanded (useful for manual
-            inspection if something goes wrong). Note that this option will
-            override all_results if True.
-
-        query_overrides : dict = {}
-            Override any part of the elasticsearch FHIR query
-
-        auth_args : Any
-            The authenication to use for the account and project (defaults to shared)
-
-        ignore_cache : bool = False
-            Bypass the caching system that auto-saves results to a CSV file.
-            Caching only occurs when all results are being retrieved.
-
-        expand_args : Any
-            Additional arguments passed to phc.Frame.expand
-
-        Examples
-        --------
-        >>> import phc.easy as phc
-        >>> phc.Auth.set({'account': '<your-account-name>'})
-        >>> phc.Project.set_current('My Project Name')
-        >>> phc.Patient.get_data_frame()
-
-        """
-        query = {
-            "type": "select",
-            "columns": "*",
-            "from": [{"table": "patient"}],
-            **query_overrides,
-        }
-
-        def transform(df: pd.DataFrame):
-            return Patient.transform_results(df, **expand_args)
-
-        return Query.execute_fhir_dsl_with_options(
-            query,
-            transform,
-            all_results,
-            raw,
-            query_overrides,
-            auth_args,
-            ignore_cache,
-        )
diff --git a/phc/easy/patients/address.py b/phc/easy/patients/address.py
index edbd1c8..c7af9a3 100644
--- a/phc/easy/patients/address.py
+++ b/phc/easy/patients/address.py
@@ -1,4 +1,6 @@
 import pandas as pd
+
+from funcy import first
 from phc.easy.codeable import generic_codeable_to_dict
 from phc.easy.util import concat_dicts
 
@@ -22,18 +24,29 @@ def expand_address_value(value):
     if type(value) is not list:
         return {}
 
-    # Value is always list of one item
-    assert len(value) == 1
-    value = value[0]
+    primary_address = first(
+        filter(lambda v: v.get("use") != "old", value)
+    ) or first(value)
+
+    other_addresses = list(
+        filter(lambda address: address != primary_address, value)
+    )
 
-    return concat_dicts(
-        [
-            expand_address_attr(f"address_{key}", item_value)
-            for key, item_value in value.items()
-            if key != "text"
-        ]
+    other_attrs = (
+        {"other_addresses": other_addresses} if len(other_addresses) > 0 else {}
     )
 
+    return {
+        **concat_dicts(
+            [
+                expand_address_attr(f"address_{key}", item_value)
+                for key, item_value in primary_address.items()
+                if key != "text"
+            ]
+        ),
+        **other_attrs,
+    }
+
 
 def expand_address_column(address_col):
     return pd.DataFrame(map(expand_address_value, address_col.values))
diff --git a/phc/easy/query/__init__.py b/phc/easy/query/__init__.py
index 45eccdf..b2bd787 100644
--- a/phc/easy/query/__init__.py
+++ b/phc/easy/query/__init__.py
@@ -69,6 +69,7 @@ class Query:
         all_results: bool = False,
         auth_args: Auth = Auth.shared(),
         callback: Union[Callable[[Any, bool], None], None] = None,
+        max_pages: Union[int, None] = None,
         log: bool = False,
         **query_kwargs,
     ):
@@ -101,6 +102,9 @@ class Query:
                     if is_finished:
                         return "batch finished
 
+        max_pages : int
+            The number of pages to retrieve (useful if working with tons of records)
+
         log : bool = False
             Whether to log the elasticsearch query sent to the server
 
@@ -151,11 +155,16 @@ class Query:
                     progress=progress,
                     callback=callback,
                     auth_args=auth_args,
+                    max_pages=max_pages,
                 ),
             )
 
         return recursive_execute_fhir_dsl(
-            query, scroll=all_results, callback=callback, auth_args=auth_args,
+            query,
+            scroll=all_results,
+            callback=callback,
+            auth_args=auth_args,
+            max_pages=max_pages,
         )
 
     @staticmethod
@@ -167,6 +176,7 @@ class Query:
         query_overrides: dict,
         auth_args: Auth,
         ignore_cache: bool,
+        max_pages: Union[int, None],
         log: bool = False,
         **query_kwargs,
     ):
@@ -179,6 +189,7 @@ class Query:
             (not ignore_cache)
             and (not raw)
             and (all_results or FhirAggregation.is_aggregation_query(query))
+            and (max_pages is None)
         )
 
         if use_cache and APICache.does_cache_for_fhir_dsl_exist(query):
@@ -191,7 +202,11 @@ class Query:
         )
 
         results = Query.execute_fhir_dsl(
-            query, all_results, auth_args, callback=callback
+            query,
+            all_results,
+            auth_args,
+            callback=callback,
+            max_pages=max_pages,
         )
 
         if isinstance(results, FhirAggregation):
diff --git a/phc/easy/query/fhir_dsl.py b/phc/easy/query/fhir_dsl.py
index 3fb03f6..b8db491 100644
--- a/phc/easy/query/fhir_dsl.py
+++ b/phc/easy/query/fhir_dsl.py
@@ -33,6 +33,8 @@ def recursive_execute_fhir_dsl(
     progress: Union[None, tqdm] = None,
     auth_args: Auth = Auth.shared(),
     callback: Union[Callable[[Any, bool], None], None] = None,
+    max_pages: Union[int, None] = None,
+    _current_page: int = 1,
     _scroll_id: str = "true",
     _prev_hits: List = [],
 ):
@@ -54,7 +56,11 @@ def recursive_execute_fhir_dsl(
     if progress:
         progress.update(current_result_count)
 
-    is_last_batch = current_result_count == 0 or scroll is False
+    is_last_batch = (
+        (current_result_count == 0)
+        or (scroll is False)
+        or ((max_pages is not None) and (_current_page >= max_pages))
+    )
     results = [] if callback else [*_prev_hits, *current_results]
 
     if callback and not is_last_batch:
@@ -73,6 +79,8 @@ def recursive_execute_fhir_dsl(
         progress=progress,
         auth_args=auth_args,
         callback=callback,
+        max_pages=max_pages,
+        _current_page=_current_page + 1,
         _scroll_id=_scroll_id,
         _prev_hits=results,
     )
diff --git a/phc/util/api_cache.py b/phc/util/api_cache.py
index 0e7f58d..475d88b 100644
--- a/phc/util/api_cache.py
+++ b/phc/util/api_cache.py
@@ -1,5 +1,6 @@
 import hashlib
 import json
+import os
 from pathlib import Path
 from typing import Callable
 
@@ -86,6 +87,9 @@ class APICache:
         writer = CSVWriter(filename)
 
         def handle_batch(batch, is_finished):
+            if is_finished and not os.path.exists(filename):
+                return pd.DataFrame()
+
             if is_finished:
                 print(f'Loading data frame from "{filename}"')
                 return APICache.read_csv(filename)


Test patch:
diff --git a/tests/test_easy_patient_address.py b/tests/test_easy_patient_address.py
index d51ff86..42bb1cd 100644
--- a/tests/test_easy_patient_address.py
+++ b/tests/test_easy_patient_address.py
@@ -1,18 +1,64 @@
 import pandas as pd
+import math
 from phc.easy.patients.address import expand_address_column
 
+
+def non_na_dict(dictionary: dict):
+    return {
+        k: v
+        for k, v in dictionary.items()
+        if not isinstance(v, float) or not math.isnan(v)
+    }
+
+
 def test_expand_address_column():
-    sample = pd.DataFrame([{
-        'address': [
-            {'line': ['123 ABC Court'], 'city': 'Zionsville', 'state': 'IN', 'use': 'home'}
+    sample = pd.DataFrame(
+        [
+            {
+                "address": [
+                    {
+                        "line": ["123 ABC Court"],
+                        "city": "Zionsville",
+                        "state": "IN",
+                        "use": "home",
+                    }
+                ]
+            },
+            {
+                "address": [
+                    {
+                        "use": "old",
+                        "state": "SC",
+                        "period": {"start": "1999", "end": "2001"},
+                    },
+                    {
+                        "state": "NC",
+                        "city": "Raleigh",
+                        "period": {"start": "2001"},
+                    },
+                ]
+            },
         ]
-    }])
+    )
 
     df = expand_address_column(sample.address)
 
-    assert df.iloc[0].to_dict() == {
-        'address_line_0': '123 ABC Court',
-        'address_city': 'Zionsville',
-        'address_state': 'IN',
-        'address_use': 'home'
+    assert non_na_dict(df.iloc[0].to_dict()) == {
+        "address_line_0": "123 ABC Court",
+        "address_city": "Zionsville",
+        "address_state": "IN",
+        "address_use": "home",
+    }
+
+    assert non_na_dict(df.iloc[1].to_dict()) == {
+        "address_city": "Raleigh",
+        "address_state": "NC",
+        "address_period_start": "2001",
+        "other_addresses": [
+            {
+                "use": "old",
+                "state": "SC",
+                "period": {"start": "1999", "end": "2001"},
+            }
+        ],
     }