Add Python Docstring

phs/data/optum.py

-"""
-
-"""
-
 import os
 import re
 import textwrap
@@ -9,9 +5,8 @@ from datetime import date
 from redivis import bigquery
 
 class Optum:
-    """
-        
-    """
+    """Build a cohort using the Optum dataset."""
+    
     # Project ID
     _project_id = 'stanfordphs'
     
@@ -266,7 +261,6 @@ class Optum:
             )
         )"""
 
-    # SQL template for getting a cohort population
     _GET_COHORT_POPULATIONS = """
         WITH {observation_period}, {person}, {condition_occurrence}, {qualifying_event}, {inclusion_criteria}, {inclusion_event},
         cohort AS (
@@ -308,19 +302,42 @@ class Optum:
     
     
     def evaluate(self, query):
-        """Returns a :obj:`pandas.DataFrame`  
-        
+        """Send a query to BigQuery to be evaluated and return as a data frame.
+         
         Args:
-        
+            query (str): A BigQuery query string
+            
         Returns:
-        
+            A data frame representing the query result
         """
         result = self._client.query(query)
         return result.to_dataframe()
 
     def get_cohort(self, zips, entry_criteria, inclusion_criteria=None, exit_criteria=None):
-        """
+        """Get a patient cohort.
+        
+        This cohort is built based on a list of entry criteria, inclusion criteria and exit
+        criteria within the given ZIP codes. 
         
+        Args:
+            zips (list): A list of US ZIP codes
+            entry_criteria (dict of str: dict): The cohort entry criteria defines the
+                conditions when people initially enter the cohort. The set of people having
+                an entry event is referred to as the initial event cohort.
+            inclusion_criteria(list, optional): The inclusion criteria defines a further
+                restriction to the initial event cohort. Each inclusion criterion will be
+                evaluated to determine the impact of the criteria on the attrition of persons
+                from the initial event cohort. The qualifying cohort is defined as all people
+                in the initial event cohort that satisfy all inclusion criteria.
+            exit_criteria(dict of str: dict, optional): The cohort exit criteria defines when
+                a person no longer qualifies for cohort membership. Cohort exit can be defined
+                in multiple ways such as the end of the observation period, a fixed time
+                interval relative to the initial entry event, the last event in a sequence of
+                related observations (e.g. persistent drug exposure) or through other
+                censoring of observation period.
+            
+        Returns:
+            A data frame representing the patient cohort.
         """
         entry_criteria_sql = self._build_entry_criteria_sql(entry_criteria)
         if inclusion_criteria:
@@ -344,10 +361,10 @@ class Optum:
                 inclusion_criteria_count=len(inclusion_criteria)))
         
         # Execute the query
-#         print(query)
         return self.evaluate(query)
 
     def _build_entry_criteria_sql(self, entry_criteria):
+        """Build query string for the entry criteria."""
         entry_criteria_sql = ""
         for k, v in entry_criteria.items():
             if k == 'event':
@@ -358,6 +375,7 @@ class Optum:
         return entry_criteria_sql
     
     def _get_condition_occurrence_parameters(self, criteria):
+        """Rebuild internally the condition occurrence parameters."""
         condition_occurrence_parameters = {}
         condition_occurrence_parameters['condition_filter'] = ""
         condition_occurrence_parameters['condition_occurrence_filter'] = ""
@@ -401,12 +419,14 @@ class Optum:
         return condition_occurrence_parameters
     
     def _build_condition_entry_criteria_sql(self, condition_occurrence_parameters):
+        """Build query string for the entry criteria based on the condition occurrence parameters."""
         return textwrap.dedent(self._ENTRY_CRITERIA.format(
                 condition_filter=condition_occurrence_parameters['condition_filter'],
                 condition_occurrence_filters=condition_occurrence_parameters['condition_occurrence_filter'],
                 continuous_observation_filters=condition_occurrence_parameters['continuous_observation_filter']))
     
     def _build_inclusion_criteria_sql(self, inclusion_criteria):
+        """Build query string for the inclusion criteria."""
         inclusion_criteria_list = []
         for index, inclusion_criterion in enumerate(inclusion_criteria, start=0):
             for k, v in inclusion_criterion.items():
@@ -442,18 +462,21 @@ class Optum:
         return ",".join(ic for ic in inclusion_criteria_list)
 
     def _build_gender_inclusion_criteria_sql(self, index, gender_filter_sql):
+        """Build query string for the inclusion criteria based on the gender parameters."""
         inclusion_criteria_sql = self._INCLUSION_CRITERIA_GENDER.format(gender_filter=gender_filter_sql)
         return self._INCLUSION_CRITERIA.format(
                 event_index=index,
                 inclusion_criteria=inclusion_criteria_sql)
 
     def _build_age_inclusion_criteria_sql(self, index, age_filter_sql):
+        """Build query string for the inclusion criteria based on the age parameters."""
         inclusion_criteria_sql = self._INCLUSION_CRITERIA_AGE.format(age_filter=age_filter_sql)
         return self._INCLUSION_CRITERIA.format(
                 event_index=index,
                 inclusion_criteria=inclusion_criteria_sql)
     
     def _build_condition_inclusion_criteria_sql(self, index, condition_occurrence_parameters):
+        """Build query string for the inclusion criteria based on the condition occurrence parameters."""
         inclusion_criteria_sql = self._INCLUSION_CRITERIA_CONDITION.format(
                 condition_filter=condition_occurrence_parameters['condition_filter'],
                 condition_occurrence_filters=condition_occurrence_parameters['condition_occurrence_filter'],
@@ -464,6 +487,7 @@ class Optum:
                 inclusion_criteria=inclusion_criteria_sql)
 
     def _build_inclusion_event_sql(self, inclusion_criteria):
+        """Build query string for the inclusion event based on the inclusion criteria."""
         ic_size = len(inclusion_criteria)
         inclusion_criteria_members = [f'SELECT inclusion_rule_id, person_id, zip FROM inclusion_event_{i}' for i in range(ic_size)]
         return self._INCLUSION_EVENT.format(