diff --git a/fieldExtraction/src/investment/investment_postprocessing_funcs.py b/fieldExtraction/src/investment/investment_postprocessing_funcs.py index e1e5eb8..ee200e7 100644 --- a/fieldExtraction/src/investment/investment_postprocessing_funcs.py +++ b/fieldExtraction/src/investment/investment_postprocessing_funcs.py @@ -79,8 +79,13 @@ def flatten_singleton_string_list(list_str: str) -> str: try: # Convert the string to a list using ast.literal_eval list_obj = ast.literal_eval(list_str) + # Check if list_obj is actually a list + if not isinstance(list_obj, list): + return list_str + # If the list has one element, return that element if len(list_obj) == 1: return list_obj[0] + # Otherwise, return the list as a string 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 9cfdb52..ca97131 100644 --- a/fieldExtraction/tests/investment_postprocess_test.py +++ b/fieldExtraction/tests/investment_postprocess_test.py @@ -40,6 +40,7 @@ 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("11"), "11") self.assertEqual(flatten_singleton_string_list("invalid"), "invalid") self.assertEqual(flatten_singleton_string_list(None), "")