PR feedback

Author: Hussain Jafari

src/pseudopeople/dataset.py

@@ -168,6 +168,19 @@ def _reformat_dates_for_noising(self) -> None:
     def drop_non_schema_columns(
         data: pd.DataFrame, dataset_schema: DatasetSchema
     ) -> pd.DataFrame:
+        """Returns data with only the columns in the dataset schema.
+
+        Parameters
+        ----------
+        data
+            The pd.DataFrame to update.
+        dataset_schema
+            A DatasetSchema which contains the columns of interest in its column attribute.
+
+        Returns
+        -------
+            A pd.DataFrame with the columns in the dataset schema.
+        """
         return data[[c.name for c in dataset_schema.columns]]
 
     @staticmethod

src/pseudopeople/noise_functions.py

@@ -216,7 +216,6 @@ def duplicate_with_guardian(
             if index_to_copy.empty:
                 continue
             noised_group_df = group_df.loc[index_to_copy]
-            noised_group_df["old_housing_type"] = noised_group_df["housing_type"]
             noised_group_df[GUARDIAN_DUPLICATION_ADDRESS_COLUMNS] = group_df.loc[
                 index_to_copy,
                 [f"{guardian}_" + column for column in GUARDIAN_DUPLICATION_ADDRESS_COLUMNS],
@@ -231,9 +230,7 @@ def duplicate_with_guardian(
         ].map(HOUSING_TYPE_GUARDIAN_DUPLICATION_RELATONSHIP_MAP)
 
         # Clean columns
-        duplicated_rows_df = duplicated_rows_df[
-            list(dataset.data.columns) + ["old_housing_type"]
-        ]
+        duplicated_rows_df = duplicated_rows_df[dataset.data.columns]
 
         # Add duplicated rows to the original data and make sure that households
         # are grouped together by sorting by date and household_id

tests/integration/release/test_release.py

@@ -241,6 +241,12 @@ def test_guardian_duplication(
     duplicated = noised.loc[noised["simulant_id"].duplicated()]
     duplicated["age"] = duplicated["age"].astype(int)
 
+    # add old housing type data to duplicated simulants
+    old_housing_data = unnoised[["simulant_id", "housing_type"]].rename(
+        {"housing_type": "unnoised_housing_type"}, axis=1
+    )
+    duplicated = duplicated.merge(old_housing_data)
+
     # separate tests for household under 18 and for college under 24
     for probability_name, age, housing_type in zip(
         [
@@ -267,7 +273,7 @@ def test_guardian_duplication(
             )
         ]
         duplicated_in_group = duplicated.loc[
-            (duplicated["age"] < age) & (duplicated["old_housing_type"] == housing_type)
+            (duplicated["age"] < age) & (duplicated["unnoised_housing_type"] == housing_type)
         ]
 
         fuzzy_checker.fuzzy_assert_proportion(