Skip to content

Commit 1ee1942

Browse files
dmorton714dmorton714
authored
Merge pull request #20 from dmorton714/4-tests-for-cleaning
Wrote a test and updated the clean salary method to catch all the edg…
2 parents 197dfc2 + 5c5345a commit 1ee1942

1 file changed

Lines changed: 253 additions & 3 deletions

File tree

‎src/notebooks/mainNb.ipynb‎

Lines changed: 253 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
},
1919
{
2020
"cell_type": "code",
21-
"execution_count": 2,
21+
"execution_count": 28,
2222
"id": "d11a2343",
2323
"metadata": {},
2424
"outputs": [],
@@ -28,7 +28,8 @@
2828
"from pathlib import Path\n",
2929
"import os\n",
3030
"import sys\n",
31-
"import re"
31+
"import re\n",
32+
"import pandas.testing as pdt"
3233
]
3334
},
3435
{
@@ -630,7 +631,8 @@
630631
},
631632
{
632633
"cell_type": "code",
633-
"execution_count": null,
634+
"execution_count": 16,
635+
634636
"id": "749ae60a",
635637
"metadata": {},
636638
"outputs": [],
@@ -782,6 +784,254 @@
782784
" return self.df"
783785
]
784786
},
787+
{
788+
"cell_type": "markdown",
789+
790+
"id": "3eb6373f",
791+
"metadata": {},
792+
"source": [
793+
"### Sample use of the clean_salary function. "
794+
]
795+
},
796+
{
797+
"cell_type": "code",
798+
"execution_count": 19,
799+
"id": "182eac4a",
800+
"metadata": {},
801+
"outputs": [
802+
{
803+
"name": "stdout",
804+
"output_type": "stream",
805+
"text": [
806+
" Salary\n",
807+
"0 50000.0\n",
808+
"1 20800.0\n",
809+
"2 104000.0\n",
810+
"3 60000.0\n",
811+
"4 75000.0\n",
812+
"5 100000.0\n",
813+
"6 150000.0\n",
814+
"7 200.0\n",
815+
"8 3000.0\n",
816+
"9 NaN\n",
817+
"10 NaN\n",
818+
"11 NaN\n",
819+
"12 145600.0\n"
820+
]
821+
}
822+
],
823+
"source": [
824+
"test_df = pd.DataFrame({\n",
825+
" \"Salary\": [\"$50k\", \"10\", \"50\", \"60,000\", \"70,000-80,000\", \"100k\", \"150000\", \"200\", \"3000\", \"5000000\", \"$1.5M\", \"invalid\", 70]\n",
826+
"})\n",
827+
"\n",
828+
"# Create instance with test DataFrame\n",
829+
"cleaner = DataCleaner(test_df)\n",
830+
"\n",
831+
"# Run salary cleaning\n",
832+
"cleaner = cleaner.clean_salary(2080)\n",
833+
"\n",
834+
"# Get the cleaned DataFrame\n",
835+
"result_df = cleaner.finalize()\n",
836+
"print(result_df)"
837+
]
838+
},
839+
{
840+
"cell_type": "code",
841+
"execution_count": 21,
842+
"id": "82806fc9",
843+
"metadata": {},
844+
"outputs": [
845+
{
846+
"name": "stdout",
847+
"output_type": "stream",
848+
"text": [
849+
" Salary\n",
850+
"0 NaN\n",
851+
"1 NaN\n",
852+
"2 NaN\n",
853+
"3 NaN\n",
854+
"4 50000.0\n",
855+
"5 5000.0\n",
856+
"6 NaN\n",
857+
"7 NaN\n",
858+
"8 NaN\n",
859+
"9 NaN\n"
860+
]
861+
}
862+
],
863+
"source": [
864+
"fail_df = pd.DataFrame({\n",
865+
" \"Salary\": [\n",
866+
" None, # NaN input\n",
867+
" \"\", # empty string\n",
868+
" \" \", # whitespace only\n",
869+
" \"abc123\", # text + numbers\n",
870+
" \"50k-abc\", # malformed range\n",
871+
" \"$-5000\", # negative salary\n",
872+
" \"∞\", # infinity symbol\n",
873+
" \"NaN\", # literal string NaN\n",
874+
" \"$1.5M\", # millions, not handled in parser\n",
875+
" \"70,000—80,000\" # em dash (—) instead of hyphen/dash\n",
876+
" ]\n",
877+
"})\n",
878+
"# Create instance with failing DataFrame\n",
879+
"fail_cleaner = DataCleaner(fail_df)\n",
880+
"# Run salary cleaning on failing DataFrame\n",
881+
"fail_cleaner = fail_cleaner.clean_salary(2080)\n",
882+
"# Get the cleaned DataFrame\n",
883+
"fail_result_df = fail_cleaner.finalize()\n",
884+
"print(fail_result_df)"
885+
]
886+
},
887+
{
888+
"cell_type": "code",
889+
"execution_count": null,
890+
"id": "123deb70",
891+
"metadata": {},
892+
"outputs": [],
893+
"source": [
894+
"class DataCleaner:\n",
895+
" def __init__(self, df: pd.DataFrame):\n",
896+
" self.df = df.copy()\n",
897+
"\n",
898+
" def clean_salary(self, hours_per_year: int = 2080):\n",
899+
" \"\"\"\n",
900+
" Clean and standardize salary values in the DataFrame.\n",
901+
"\n",
902+
" Steps performed:\n",
903+
" 1. Remove currency symbols, commas, and shorthand (e.g., \"$50k\" → 50000).\n",
904+
" 2. Handle ranges by converting them to the average value \n",
905+
" (e.g., \"50,000–70,000\" → 60000).\n",
906+
" 3. Handle shorthand \"M\" (e.g., \"$1.5M\" → 1,500,000).\n",
907+
" 4. Convert values to numeric, coercing invalid entries to NaN.\n",
908+
" 5. Treat values <= 200 as hourly wages and convert to annual salaries \n",
909+
" (multiplied by `hours_per_year`).\n",
910+
" 6. Drop unrealistic values greater than 1,000,000 (set to NaN).\n",
911+
"\n",
912+
" Parameters\n",
913+
" ----------\n",
914+
" hours_per_year : int, optional (default=2080)\n",
915+
" Number of work hours in a year for converting hourly to annual salary.\n",
916+
"\n",
917+
" Returns\n",
918+
" -------\n",
919+
" self : object\n",
920+
" The current instance with the cleaned Salary column.\n",
921+
" \"\"\"\n",
922+
" try:\n",
923+
" if \"Salary\" in self.df.columns:\n",
924+
" self.df[\"Salary\"] = self.df[\"Salary\"].astype(str)\n",
925+
"\n",
926+
" def parse_salary(val: str):\n",
927+
" val = val.strip()\n",
928+
" if not val or val.lower() in {\"nan\", \"none\"}:\n",
929+
" return None\n",
930+
"\n",
931+
" # Normalize dash types (hyphen, en dash, em dash \"-\")\n",
932+
" val = re.sub(r\"[–—]\", \"-\", val)\n",
933+
"\n",
934+
" # Handle range like \"50k-70k\" or \"50,000-70,000\"\n",
935+
" if \"-\" in val:\n",
936+
" parts = val.split(\"-\")\n",
937+
" nums = [parse_salary(p) for p in parts if p.strip()]\n",
938+
" nums = [n for n in nums if n is not None]\n",
939+
" return sum(nums) / len(nums) if nums else None\n",
940+
"\n",
941+
" # Remove $, commas, spaces\n",
942+
" val = re.sub(r\"[\\$,]\", \"\", val)\n",
943+
"\n",
944+
" # Handle shorthand k/K (e.g., \"50k\" → 50000)\n",
945+
" match_k = re.match(r\"^(\\d+(\\.\\d+)?)[kK]$\", val)\n",
946+
" if match_k:\n",
947+
" return float(match_k.group(1)) * 1000\n",
948+
"\n",
949+
" # Handle shorthand M (e.g., \"1.5M\" → 1500000)\n",
950+
" match_m = re.match(r\"^(\\d+(\\.\\d+)?)[mM]$\", val)\n",
951+
" if match_m:\n",
952+
" return float(match_m.group(1)) * 1_000_000\n",
953+
"\n",
954+
" # Plain number (integer or float)\n",
955+
" try:\n",
956+
" return float(val)\n",
957+
" except ValueError:\n",
958+
" return None\n",
959+
"\n",
960+
" # Apply parsing\n",
961+
" self.df[\"Salary\"] = self.df[\"Salary\"].apply(parse_salary)\n",
962+
"\n",
963+
" # Convert small numbers (hourly) to annual\n",
964+
" self.df.loc[self.df[\"Salary\"] <= 200, \"Salary\"] *= hours_per_year\n",
965+
"\n",
966+
" # Drop unrealistic salaries\n",
967+
" self.df.loc[self.df[\"Salary\"] > 1_000_000, \"Salary\"] = None\n",
968+
"\n",
969+
" except Exception as e:\n",
970+
" print(f\"[Warning] Failed salary cleaning: {e}\")\n",
971+
"\n",
972+
" return self\n",
973+
"\n",
974+
" def finalize(self):\n",
975+
" \"\"\"Return cleaned dataframe.\"\"\"\n",
976+
" return self.df\n"
977+
]
978+
},
979+
{
980+
"cell_type": "code",
981+
"execution_count": null,
982+
"id": "688bdf74",
983+
"metadata": {},
984+
"outputs": [
985+
{
986+
"name": "stdout",
987+
"output_type": "stream",
988+
"text": [
989+
"✅ Salary cleaning DataFrame test passed!\n"
990+
]
991+
}
992+
],
993+
"source": [
994+
"# Test DataFrame with edge/fail cases\n",
995+
"fail_df = pd.DataFrame({\n",
996+
" \"Salary\": [\n",
997+
" None, # NaN\n",
998+
" \"\", # NaN\n",
999+
" \" \", # NaN\n",
1000+
" \"abc123\", # NaN\n",
1001+
" \"50k-abc\", # 50000.0\n",
1002+
" \"$-5000\", # -5000.0 (still allowed for now)\n",
1003+
" \"∞\", # NaN\n",
1004+
" \"NaN\", # NaN\n",
1005+
" \"$1.5M\", # NaN ( >1,000,000 rule)\n",
1006+
" \"70,000—80,000\" # 75000.0 (dash normalized)\n",
1007+
" ]\n",
1008+
"})\n",
1009+
"\n",
1010+
"# Run through cleaner\n",
1011+
"cleaner = DataCleaner(fail_df)\n",
1012+
"result = cleaner.clean_salary().finalize().reset_index(drop=True)\n",
1013+
"\n",
1014+
"# Expected results as DataFrame\n",
1015+
"expected = pd.DataFrame({\n",
1016+
" \"Salary\": [\n",
1017+
" None, # None\n",
1018+
" None, # empty string\n",
1019+
" None, # whitespace\n",
1020+
" None, # abc123\n",
1021+
" 50000.0, # 50k-abc\n",
1022+
" 5000.0, # negative salary\n",
1023+
" None, # infinity\n",
1024+
" None, # \"NaN\"\n",
1025+
" None, # 1.5M filtered out\n",
1026+
" 75000.0 # range with em dash\n",
1027+
" ]\n",
1028+
"}, dtype=\"float64\").reset_index(drop=True)\n",
1029+
"\n",
1030+
"# Assertion test\n",
1031+
"pdt.assert_frame_equal(result, expected)\n",
1032+
"print(\"✅ Salary cleaning DataFrame test passed!\")"
1033+
]
1034+
},
7851035
{
7861036
"cell_type": "markdown",
7871037
"id": "6ddbb4c0",

0 commit comments

Comments
 (0)