Code tidy

phs/data/optum.py

@@ -272,7 +272,7 @@ class Optum:
                     q.person_id,
                     q.zip,
                     q.index_date,
-                    SUM(COALESCE(CAST(POWER(CAST(2 as int64), i.inclusion_rule_id) AS int64), 0)) as inclusion_rule_mask
+                    SUM(COALESCE(CAST(POWER(CAST(2 as int64), i.inclusion_rule_id) AS int64), 0)) AS inclusion_rule_mask
                 FROM qualifying_event q
                 LEFT JOIN inclusion_event i 
                 ON i.person_id = q.person_id AND i.zip = q.zip
@@ -370,8 +370,8 @@ class Optum:
             if k == 'event':
                 event_type = list(v.keys())[0]
                 if event_type == 'condition':
-                    condition_occurrence_parameters = self._get_condition_occurrence_parameters(entry_criteria)
-                    entry_criteria_sql = self._build_condition_entry_criteria_sql(condition_occurrence_parameters)
+                    params = self._get_condition_occurrence_parameters(entry_criteria)
+                    entry_criteria_sql = self._build_condition_entry_criteria_sql(params)
         return entry_criteria_sql
     
     def _get_condition_occurrence_parameters(self, criteria):
@@ -380,42 +380,59 @@ class Optum:
         condition_occurrence_parameters['condition_filter'] = ""
         condition_occurrence_parameters['condition_occurrence_filter'] = ""
         condition_occurrence_parameters['condition_restriction_before_index_date_filter'] = ""
-        condition_occurrence_parameters['condition_restriction_after_index_date_filter'] = "AND e.condition_start_date <= DATE_ADD(CAST(p.index_date AS date), INTERVAL 0 DAY)"
+        condition_occurrence_parameters['condition_restriction_after_index_date_filter'] = \
+                "AND e.condition_start_date <= DATE_ADD(CAST(p.index_date AS date), INTERVAL 0 DAY)"
         condition_occurrence_parameters['continuous_observation_filter'] = ""
         for k, v in criteria.items():
             if k == 'event':
                 condition_value = v['condition']
                 if type(condition_value) is str:
-                    condition_valueset = ",".join(f'"{x}"' for x in self._CONDITIONS[condition_value]['codes'])
-                    condition_occurrence_parameters['condition_filter'] = f'WHERE condition_source_value IN ({condition_valueset})'
+                    codes = self._CONDITIONS[condition_value]['codes']
+                    condition_valueset = ",".join(f'"{x}"' for x in codes)
+                    condition_occurrence_parameters['condition_filter'] = \
+                            f'WHERE condition_source_value IN ({condition_valueset})'
                 elif type(condition_value) is list:
                     condition_valueset = ",".join(f'"{x}"' for x in condition_value)
-                    condition_occurrence_parameters['condition_filter'] = f'WHERE condition_source_value IN ({condition_valueset})'
+                    condition_occurrence_parameters['condition_filter'] = \
+                            f'WHERE condition_source_value IN ({condition_valueset})'
             elif k == 'occurrence':
                 occurrence_type = list(v.keys())[0]
                 if occurrence_type == 'startDateBefore':
                     occurrence_date = v['startDateBefore']
-                    condition_occurrence_parameters['condition_occurrence_filter'] = f'WHERE c.ordinal = 1 AND c.condition_start_date < "{occurrence_date}"'
+                    condition_occurrence_parameters['condition_occurrence_filter'] = \
+                            f'WHERE c.ordinal = 1 AND c.condition_start_date < "{occurrence_date}"'
                 elif occurrence_type == 'startDateAfter':
                     occurrence_date = v['startDateAfter']
-                    condition_occurrence_parameters['condition_occurrence_filter'] = f'WHERE c.ordinal = 1 AND c.condition_start_date > "{occurrence_date}"'
+                    condition_occurrence_parameters['condition_occurrence_filter'] = \
+                            f'WHERE c.ordinal = 1 AND c.condition_start_date > "{occurrence_date}"'
                 elif occurrence_type == 'startDateAt':
                     occurrence_date = v['startDateAt']
-                    condition_occurrence_parameters['condition_occurrence_filter'] = f'WHERE c.ordinal = 1 AND c.condition_start_date = "{occurrence_date}"'
+                    condition_occurrence_parameters['condition_occurrence_filter'] = \
+                            f'WHERE c.ordinal = 1 AND c.condition_start_date = "{occurrence_date}"'
                 elif occurrence_type == 'startDateBetween':
                     occurrence_start_date = v['startDateBetween'].split('/')[0]
                     occurrence_end_date = v['startDateBetween'].split('/')[1]
-                    condition_occurrence_parameters['condition_occurrence_filter'] = f'WHERE c.ordinal = 1 AND c.condition_start_date >= "{occurrence_start_date}" AND c.condition_start_date <= "{occurrence_end_date}"'
+                    condition_occurrence_parameters['condition_occurrence_filter'] = \
+                            f'WHERE c.ordinal = 1 AND c.condition_start_date >= "{occurrence_start_date}" ' \
+                            f'AND c.condition_start_date <= "{occurrence_end_date}"'
             elif k == 'restriction':
                 for restriction_type, restriction_value in v.items():
                     if restriction_type == 'withinDaysBeforeIndexDate':
-                        condition_occurrence_parameters['condition_restriction_before_index_date_filter'] = f'AND e.condition_start_date >= DATE_ADD(CAST(p.index_date AS date), INTERVAL -{restriction_value} DAY)'
+                        condition_occurrence_parameters['condition_restriction_before_index_date_filter'] = \
+                                f'AND e.condition_start_date >= DATE_ADD(CAST(p.index_date AS date), ' \
+                                f'INTERVAL -{restriction_value} DAY)'
                     elif restriction_type == 'withinDaysAfterIndexDate':
-                        condition_occurrence_parameters['condition_restriction_after_index_date_filter'] = f'AND e.condition_start_date <= DATE_ADD(CAST(p.index_date AS date), INTERVAL {restriction_value} DAY)'
+                        condition_occurrence_parameters['condition_restriction_after_index_date_filter'] = \
+                                f'AND e.condition_start_date <= DATE_ADD(CAST(p.index_date AS date), ' \
+                                f'INTERVAL {restriction_value} DAY)'
             elif k == 'continuousObservation':
                 days_before = v['minDaysBeforeIndexDate'] if 'minDaysBeforeIndexDate' in v else 0
                 days_after = v['minDaysAfterIndexDate'] if 'minDaysAfterEntryEvent' in v else 0
-                condition_occurrence_parameters['continuous_observation_filter'] = f'WHERE DATE_ADD(op.observation_period_start_date, interval {days_before} DAY) <= e.condition_start_date AND DATE_ADD(e.condition_start_date, interval {days_after} DAY) <= op.observation_period_end_date'         
+                condition_occurrence_parameters['continuous_observation_filter'] = \
+                        f'WHERE DATE_ADD(op.observation_period_start_date, ' \
+                        f'INTERVAL {days_before} DAY) <= e.condition_start_date ' \
+                        f'AND DATE_ADD(e.condition_start_date, ' \
+                        f'INTERVAL {days_after} DAY) <= op.observation_period_end_date'         
         return condition_occurrence_parameters
     
     def _build_condition_entry_criteria_sql(self, condition_occurrence_parameters):
@@ -438,14 +455,17 @@ class Optum:
                 elif k == 'ageGroup':
                     start_age = self._AGE_GROUPS[v]['start']
                     end_age = self._AGE_GROUPS[v]['end']
-                    filter_sql = f'WHERE EXTRACT(YEAR FROM index_date)-year_of_birth >= {start_age} AND EXTRACT(YEAR FROM index_date)-year_of_birth <= {end_age}'
+                    filter_sql = f'WHERE EXTRACT(YEAR FROM index_date)-year_of_birth >= {start_age} ' \
+                            f'AND EXTRACT(YEAR FROM index_date)-year_of_birth <= {end_age}'
                     inclusion_criteria_sql = self._build_age_inclusion_criteria_sql(index, filter_sql)
                     inclusion_criteria_list.append(inclusion_criteria_sql)
                 elif k == 'ageRange':
-                    if not re.match(r"[0-9]*-[0-9]*", v): raise ValueError("Unable to parse value range expression. Some valid examples: -22, or 16-22, or 22-")
+                    if not re.match(r"[0-9]*-[0-9]*", v):
+                        raise ValueError("Unable to parse value range expression. Some valid examples: -22, or 16-22, or 22-")
                     start_age = int(v.split('-')[0]) if v.split('-')[0] != '' else 0
                     end_age = int(v.split('-')[1]) if v.split('-')[1] != '' else 999
-                    filter_sql = f'WHERE EXTRACT(YEAR FROM index_date)-year_of_birth >= {start_age} AND EXTRACT(YEAR FROM index_date)-year_of_birth <= {end_age}'
+                    filter_sql = f'WHERE EXTRACT(YEAR FROM index_date)-year_of_birth >= {start_age} ' \
+                            f'AND EXTRACT(YEAR FROM index_date)-year_of_birth <= {end_age}'
                     inclusion_criteria_sql = self._build_age_inclusion_criteria_sql(index, filter_sql)
                     inclusion_criteria_list.append(inclusion_criteria_sql)
                 elif k == 'ageAt':