diff --git a/fieldExtraction/src/investment/investment_postprocessing_funcs.py b/fieldExtraction/src/investment/investment_postprocessing_funcs.py index effbd5d..fe4a046 100644 --- a/fieldExtraction/src/investment/investment_postprocessing_funcs.py +++ b/fieldExtraction/src/investment/investment_postprocessing_funcs.py @@ -74,12 +74,14 @@ def flatten_singleton_string_list(list_str: str) -> str: Returns: str: The single element if the list has one element, otherwise the list as a string. """ + if string_utils.is_empty(list_str): + return "" try: # Convert the string to a list using ast.literal_eval list_obj = ast.literal_eval(list_str) if len(list_obj) == 1: return list_obj[0] - return list_obj + return ", ".join(list_obj) except (ValueError, SyntaxError): return list_str diff --git a/fieldExtraction/tests/investment_postprocess_test.py b/fieldExtraction/tests/investment_postprocess_test.py index d598f7d..9cfdb52 100644 --- a/fieldExtraction/tests/investment_postprocess_test.py +++ b/fieldExtraction/tests/investment_postprocess_test.py @@ -25,10 +25,11 @@ class TestPostprocessFunctions(unittest.TestCase): self.assertEqual(normalize_indicator_field(None), "N") def test_format_rate_fields_with_commas(self): + # Test numeric values self.assertEqual(format_rate_fields_with_commas("1234.567"), "1,234.57") self.assertEqual(format_rate_fields_with_commas("1000"), "1,000.00") - self.assertEqual(format_rate_fields_with_commas("abc"), "0.00") - self.assertEqual(format_rate_fields_with_commas(None), "0.00") + self.assertEqual(format_rate_fields_with_commas(1234.567), "1,234.57") + self.assertEqual(format_rate_fields_with_commas(1000), "1,000.00") def test_remove_hyphens(self): self.assertEqual(remove_hyphens("123-45-6789"), "123456789") @@ -38,9 +39,9 @@ class TestPostprocessFunctions(unittest.TestCase): def test_flatten_singleton_string_list(self): self.assertEqual(flatten_singleton_string_list("['123']"), "123") - self.assertEqual(flatten_singleton_string_list("['123', '456']"), ['123', '456']) + self.assertEqual(flatten_singleton_string_list("['123', '456']"), "123, 456") self.assertEqual(flatten_singleton_string_list("invalid"), "invalid") - self.assertEqual(flatten_singleton_string_list(None), None) + self.assertEqual(flatten_singleton_string_list(None), "") @patch("src.utils.llm_utils.invoke_claude")