diff --git a/src/CountriesMerger.py b/src/CountriesMerger.py index a406f310680..dd17eed83d0 100644 --- a/src/CountriesMerger.py +++ b/src/CountriesMerger.py @@ -7,7 +7,7 @@ class CountriesMerger: def mergeSrcIntoDst(src: pd.Series, dst: pd.Series): def merge(series): series = series.dropna() - return sorted(series[0].union(series[1:])) + return sorted(set().union(*series)) return (pd .merge(dst, src, how='left', left_index=True, right_index=True) .apply(merge, axis='columns')) diff --git a/src/CountriesMergerTest.py b/src/CountriesMergerTest.py index 89330cd4c39..a5a7a2272e2 100644 --- a/src/CountriesMergerTest.py +++ b/src/CountriesMergerTest.py @@ -8,8 +8,8 @@ class CountriesMergerTest(unittest.TestCase): def test_mergeSrcIntoDst(self): # Given - dstCountries = pd.Series({'NO72A': {'CountryA-1'}, 'EW096': {'CountryA-2'}}, name='dstCountries') - srcCountries = pd.Series({'NO72A': 'CountryB-1', 'FS192': 'CountryB-2'}, name='srcCountries') + dstCountries = pd.Series({'NO72A': ['CountryA-1'], 'EW096': ['CountryA-2']}, name='dstCountries') + srcCountries = pd.Series({'NO72A': ['CountryB-1'], 'FS192': ['CountryB-2']}, name='srcCountries') # When mergedCountries = CountriesMerger.mergeSrcIntoDst(src = srcCountries, dst = dstCountries)