refining CountryColumnAdderTest

src/CountryColumnAdder.py

@@ -1,24 +1,35 @@
 import pycountry
+import pandas as pd
 
 class CountryColumnAdder:
     
-    @staticmethod
+    def __init__(self, dataFrame_SPLTTYPE_By_VAERS_ID):
-    def addCountryColumn(dataFrame):
+        self.dataFrame_COUNTRY_By_VAERS_ID = self._create_dataFrame_COUNTRY_By_VAERS_ID(dataFrame_SPLTTYPE_By_VAERS_ID)
-        dataFrame['COUNTRY'] = CountryColumnAdder._splttype2Country(dataFrame['SPLTTYPE'])
-        return dataFrame
         
-    @staticmethod
+    def addCountryColumn(self, dataFrame):
-    def _splttype2Country(splttypeSeries):
+        return pd.merge(
+            dataFrame,
+            self.dataFrame_COUNTRY_By_VAERS_ID,
+            how = 'left',
+            left_index = True,
+            right_index = True)
+
+    def _create_dataFrame_COUNTRY_By_VAERS_ID(self, dataFrame_SPLTTYPE_By_VAERS_ID):
+        dataFrame_COUNTRY_By_VAERS_ID = dataFrame_SPLTTYPE_By_VAERS_ID[['SPLTTYPE']].copy()
+        dataFrame_COUNTRY_By_VAERS_ID['COUNTRY'] = self._splttype2Country(dataFrame_COUNTRY_By_VAERS_ID['SPLTTYPE'])
+        dataFrame_COUNTRY_By_VAERS_ID = dataFrame_COUNTRY_By_VAERS_ID.drop(columns = ['SPLTTYPE'])
+        return dataFrame_COUNTRY_By_VAERS_ID
+
+    def _splttype2Country(self, splttypeSeries):
         return (splttypeSeries
                 .apply(
                     lambda splttype:
-                        CountryColumnAdder._getCountryNameOfSplttypeOrDefault(
+                        self._getCountryNameOfSplttypeOrDefault(
                             splttype = splttype,
                             default = 'Unknown Country'))
                 .astype("string"))
 
-    @staticmethod
+    def _getCountryNameOfSplttypeOrDefault(self, splttype, default):
-    def _getCountryNameOfSplttypeOrDefault(splttype, default):
         if not isinstance(splttype, str):
             return default
         

src/CountryColumnAdderTest.py

@@ -14,13 +14,16 @@ class CountryColumnAdderTest(unittest.TestCase):
             data = [  ['GBPFIZER INC2020486806'],
                       ['FRMODERNATX, INC.MOD20224'],
                       ['dummy']],
-            index = [
+            index = pd.Index(
-                "4711",
+                    name = 'VAERS_ID',
-                "0815",
+                    data = [
-                "123"])
+                        "4711",
+                        "0815",
+                        "123"]))
+        countryColumnAdder = CountryColumnAdder(dataFrame)
 
         # When
-        dataFrameWithCountryColumn = CountryColumnAdder.addCountryColumn(dataFrame)
+        dataFrameWithCountryColumn = countryColumnAdder.addCountryColumn(dataFrame)
 
         # Then
         assert_frame_equal(
@@ -30,8 +33,56 @@ class CountryColumnAdderTest(unittest.TestCase):
             data = [  ['GBPFIZER INC2020486806',    'United Kingdom'],
                       ['FRMODERNATX, INC.MOD20224', 'France'],
                       ['dummy',                     'Unknown Country']],
-            index = [
+            index = pd.Index(
-                "4711",
+                    name = 'VAERS_ID',
-                "0815",
+                    data = [
-                "123"],
+                        "4711",
+                        "0815",
+                        "123"]),
+            dtypes = {'COUNTRY': 'string'}))
+
+
+    def test_addCountryColumn2(self):
+        # Given
+        countryColumnAdder = CountryColumnAdder(
+            TestHelper.createDataFrame(
+                columns = ['SPLTTYPE'],
+                data = [  ['GBPFIZER INC2020486806'],
+                          ['FRMODERNATX, INC.MOD20224'],
+                          ['dummy']],
+                index = pd.Index(
+                    name = 'VAERS_ID',
+                    data = [
+                        2547744,
+                        2547730,
+                        2540815])))
+        dataFrame = TestHelper.createDataFrame(
+            columns = ['VAX_LOT'],
+            data = [  ['1808982'],
+                      ['EW0175'],
+                      ['EW0176']],
+            index = pd.Index(
+                    name = 'VAERS_ID',
+                    data = [
+                        2547730,
+                        2547730,
+                        2547744]))
+        
+        # When
+        dataFrameWithCountryColumn = countryColumnAdder.addCountryColumn(dataFrame)
+
+        # Then
+        assert_frame_equal(
+            dataFrameWithCountryColumn,
+            TestHelper.createDataFrame(
+            columns = ['VAX_LOT', 'COUNTRY'],
+            data = [  ['1808982', 'France'],
+                      ['EW0175',  'France'],
+                      ['EW0176',  'United Kingdom']],
+            index = pd.Index(
+                    name = 'VAERS_ID',
+                    data = [
+                        2547730,
+                        2547730,
+                        2547744]),
             dtypes = {'COUNTRY': 'string'}))

src/VaersReader.py

@@ -16,7 +16,7 @@ def getVaersForYears(years):
 def getNonDomesticVaers():
     return _getVaers(
         [_getVaersDescrReader().readNonDomesticVaersDescr()],
-        CountryColumnAdder.addCountryColumn)
+        addCountryColumn = lambda dataFrame: CountryColumnAdder(dataFrame).addCountryColumn(dataFrame))
 
 def _getVaersDescrReader():
     return VaersDescrReader(dataDir = "VAERS")