Changing self.office to self.office_id in BaseDataHandler and subclasses

washingtonpost · Dec 12, 2024 · f2ca6cf · f2ca6cf
1 parent ec5e072
commit f2ca6cf
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 9 deletions.
diff --git a/src/elexmodel/handlers/data/BaseDataHandler.py b/src/elexmodel/handlers/data/BaseDataHandler.py
@@ -14,10 +14,10 @@ class BaseDataHandler(abc.ABC):
     """
 
     def __init__(
-        self, election_id, office, geographic_unit_type, estimands, s3_client=None, historical=False, data=None
+        self, election_id, office_id, geographic_unit_type, estimands, s3_client=None, historical=False, data=None
     ):
         self.election_id = election_id
-        self.office = office
+        self.office_id = office_id
         self.geographic_unit_type = geographic_unit_type
         self.estimands = estimands
         self.s3_client = s3_client
@@ -32,15 +32,15 @@ def __init__(
 
     def get_data_path(self):
         directory_path = get_directory_path()
-        path = f"{directory_path}/data/{self.election_id}/{self.office}/data_{self.geographic_unit_type}.csv"
+        path = f"{directory_path}/data/{self.election_id}/{self.office_id}/data_{self.geographic_unit_type}.csv"
         return path
 
     def get_data(self):
         # If local data file is not available, read data from s3
         if not Path(self.file_path).is_file():
             path_info = {
                 "election_id": self.election_id,
-                "office": self.office,
+                "office": self.office_id,
                 "geographic_unit_type": self.geographic_unit_type,
             }
             file_path = self.s3_client.get_file_path("preprocessed", path_info)

diff --git a/src/elexmodel/handlers/data/PreprocessedData.py b/src/elexmodel/handlers/data/PreprocessedData.py
@@ -15,7 +15,7 @@ class PreprocessedDataHandler(BaseDataHandler):
     def __init__(
         self,
         election_id,
-        office,
+        office_id,
         geographic_unit_type,
         estimands,
         estimand_baselines,
@@ -30,7 +30,13 @@ def __init__(
         self.estimand_baselines = estimand_baselines
         self.include_results_estimand = include_results_estimand
         super().__init__(
-            election_id, office, geographic_unit_type, estimands, s3_client=s3_client, historical=historical, data=data
+            election_id,
+            office_id,
+            geographic_unit_type,
+            estimands,
+            s3_client=s3_client,
+            historical=historical,
+            data=data,
         )
 
     def select_rows_in_states(self, data, states_with_election):
@@ -45,7 +51,7 @@ def load_data(self, data):
         """
         Load preprocessed csv data as df
         """
-        LOG.info("Loading preprocessed data: %s, %s, %s", self.election_id, self.office, self.geographic_unit_type)
+        LOG.info("Loading preprocessed data: %s, %s, %s", self.election_id, self.office_id, self.geographic_unit_type)
         data = self.estimandizer.add_estimand_baselines(
             data,
             self.estimand_baselines,

diff --git a/src/elexmodel/handlers/data/VersionedData.py b/src/elexmodel/handlers/data/VersionedData.py
@@ -65,7 +65,7 @@ def get_versioned_results(self, filepath=None):
         if self.election_id.startswith("2020-11-03_USA_G"):
             path = "elex-models-prod/2020-general/results/pres/current.csv"
         else:
-            base_dir = f"{S3_FILE_PATH}/{self.election_id}/results/{self.office}/{self.geographic_unit_type}"
+            base_dir = f"{S3_FILE_PATH}/{self.election_id}/results/{self.office_id}/{self.geographic_unit_type}"
             if self.election_id.startswith("2024-11-05_USA_G"):
                 path = base_dir + "/current_counties.csv"
             else:
@@ -237,6 +237,6 @@ def get_versioned_predictions(self, filepath=None):
         if self.election_id.startswith("2020-11-03_USA_G"):
             raise ValueError("No versioned predictions available for this election.")
 
-        path = f"{S3_FILE_PATH}/{self.election_id}/predictions/{self.office}/{self.geographic_unit_type}/current.csv"
+        path = f"{S3_FILE_PATH}/{self.election_id}/predictions/{self.office_id}/{self.geographic_unit_type}/current.csv"
 
         return self.s3_client.get(path, self.sample)