diff --git a/src/parent_child/pipeline.py b/src/parent_child/pipeline.py index 441adcc..ba471eb 100644 --- a/src/parent_child/pipeline.py +++ b/src/parent_child/pipeline.py @@ -808,7 +808,11 @@ def assign_child_ranks(df, grouper_col="grouping_key"): try: amendment_val = int(float(raw)) except (ValueError, TypeError): - amendment_val = raw + raw_str = str(raw).strip() + if len(raw_str) == 1 and raw_str.isalpha(): + amendment_val = ord(raw_str.upper()) - ord("A") + 1 + else: + amendment_val = raw rank_val = f"{parent_rank_str}.{amendment_val}.{i}" children.loc[idx, "child_rank"] = rank_val diff --git a/src/parent_child/qc.py b/src/parent_child/qc.py index fa95510..ea81c20 100644 --- a/src/parent_child/qc.py +++ b/src/parent_child/qc.py @@ -2032,7 +2032,10 @@ class ParentChildEngine: try: return str(int(float(x))) except (ValueError, TypeError): - return str(x).strip() + raw_str = str(x).strip() + if len(raw_str) == 1 and raw_str.isalpha(): + return str(ord(raw_str.upper()) - ord("A") + 1) + return raw_str return "0" def prepare_output(self, df: pd.DataFrame) -> pd.DataFrame: diff --git a/src/tests/test_parent_child.py b/src/tests/test_parent_child.py index 48e27f0..59ebaa3 100644 --- a/src/tests/test_parent_child.py +++ b/src/tests/test_parent_child.py @@ -776,7 +776,7 @@ class TestAssignChildRanksWithAmendmentNum(unittest.TestCase): self.assertEqual(parent_rank, "1") def test_children_use_character_amendment_num_in_rank(self): - """Children with character amendment number (A, B, C) use it in rank.""" + """Children with character amendment number (A, B, C) convert to numeric (A=1, B=2, etc.).""" df = pd.DataFrame( { "grouping_key": ["TIN:123", "TIN:123", "TIN:123"], @@ -791,11 +791,11 @@ class TestAssignChildRanksWithAmendmentNum(unittest.TestCase): child_df = result[result["parent"] == False].sort_values("fixed_effective_date") ranks = child_df["child_rank"].tolist() - self.assertEqual(ranks[0], "1.A.1") # Amendment A, sequence 1 - self.assertEqual(ranks[1], "1.B.2") # Amendment B, sequence 2 + self.assertEqual(ranks[0], "1.1.1") # Amendment A=1, sequence 1 + self.assertEqual(ranks[1], "1.2.2") # Amendment B=2, sequence 2 def test_orphans_use_character_amendment_num_in_rank(self): - """Orphans with character amendment number get 0.amendment.seq format.""" + """Orphans with character amendment number convert to numeric (A=1, C=3).""" df = pd.DataFrame( { "grouping_key": ["TIN:123", "TIN:123"], @@ -808,8 +808,8 @@ class TestAssignChildRanksWithAmendmentNum(unittest.TestCase): result = assign_child_ranks(df) ranks = result["child_rank"].tolist() - self.assertEqual(ranks[0], "0.A.1") - self.assertEqual(ranks[1], "0.C.2") + self.assertEqual(ranks[0], "0.1.1") # A=1 + self.assertEqual(ranks[1], "0.3.2") # C=3 def test_mixed_numeric_and_character_amendment_nums(self): """Mix of numeric and character amendment numbers in same group.""" @@ -833,7 +833,7 @@ class TestAssignChildRanksWithAmendmentNum(unittest.TestCase): child_df = result[result["parent"] == False].sort_values("fixed_effective_date") ranks = child_df["child_rank"].tolist() self.assertEqual(ranks[0], "1.3.1") # numeric 3 - self.assertEqual(ranks[1], "1.A.2") # character A + self.assertEqual(ranks[1], "1.1.2") # character A=1 self.assertEqual(ranks[2], "1.5.3") # numeric 5 @@ -962,7 +962,7 @@ class TestComputeRanksWithAmendmentNum(unittest.TestCase): ) def test_child_rank_dest_with_character_amendment_num(self): - """child_rank_dest handles character amendment values (A, B, C).""" + """child_rank_dest converts character amendment values (A=1, B=2, C=3).""" engine = self._make_engine() df = self._make_base_df( extra_cols={ @@ -971,13 +971,13 @@ class TestComputeRanksWithAmendmentNum(unittest.TestCase): ) result = engine.compute_ranks(df) - # Index 2 (eff 2024-02-01, amendment=A) comes first -> seq 1 - # Index 1 (eff 2024-03-01, amendment=B) comes second -> seq 2 - self.assertEqual(result.at[2, "child_rank_dest"], "1.A.1") - self.assertEqual(result.at[1, "child_rank_dest"], "1.B.2") + # Index 2 (eff 2024-02-01, amendment=A=1) comes first -> seq 1 + # Index 1 (eff 2024-03-01, amendment=B=2) comes second -> seq 2 + self.assertEqual(result.at[2, "child_rank_dest"], "1.1.1") + self.assertEqual(result.at[1, "child_rank_dest"], "1.2.2") def test_orphan_rank_dest_with_character_amendment_num(self): - """orphan_rank_dest handles character amendment values.""" + """orphan_rank_dest converts character amendment values (A=1).""" engine = self._make_engine() df = self._make_base_df( extra_cols={ @@ -986,7 +986,7 @@ class TestComputeRanksWithAmendmentNum(unittest.TestCase): ) result = engine.compute_ranks(df) - self.assertEqual(result.at[3, "orphan_rank_dest"], "0.A.1") + self.assertEqual(result.at[3, "orphan_rank_dest"], "0.1.1") # A=1 def test_child_rank_dest_with_mixed_numeric_and_character_amendment(self): """child_rank_dest handles mix of numeric and character amendment values.""" @@ -999,11 +999,11 @@ class TestComputeRanksWithAmendmentNum(unittest.TestCase): result = engine.compute_ranks(df) # Index 2 (eff 2024-02-01, amendment=3) -> seq 1 - # Index 1 (eff 2024-03-01, amendment=A) -> seq 2 + # Index 1 (eff 2024-03-01, amendment=A=1) -> seq 2 self.assertEqual(result.at[2, "child_rank_dest"], "1.3.1") - self.assertEqual(result.at[1, "child_rank_dest"], "1.A.2") - # Index 3 is orphan with amendment=B - self.assertEqual(result.at[3, "orphan_rank_dest"], "0.B.1") + self.assertEqual(result.at[1, "child_rank_dest"], "1.1.2") # A=1 + # Index 3 is orphan with amendment=B=2 + self.assertEqual(result.at[3, "orphan_rank_dest"], "0.2.1") # B=2 if __name__ == "__main__":