refactoring

src/BatchCodeTableFactory.py

@@ -0,0 +1,45 @@
+import pandas as pd
+from CompanyColumnAdder import CompanyColumnAdder
+from SummationTableFactory import SummationTableFactory
+
+class BatchCodeTableFactory:
+
+    def __init__(self, dataFrame: pd.DataFrame):
+        self.dataFrame = dataFrame
+        self.companyColumnAdder = CompanyColumnAdder(dataFrame)
+        self.countryBatchCodeTable = SummationTableFactory.createSummationTable(
+            dataFrame.groupby(
+                [
+                    dataFrame['COUNTRY'],
+                    dataFrame['VAX_LOT']
+                ]))
+
+    def createGlobalBatchCodeTable(self):
+        return self._postProcess(SummationTableFactory.createSummationTable(self.dataFrame.groupby('VAX_LOT')))
+
+    def createBatchCodeTableByCountry(self, country):
+        return self._postProcess(self._getBatchCodeTableByCountry(country))
+
+    def _postProcess(self, batchCodeTable):
+        batchCodeTable = self.companyColumnAdder.addCompanyColumn(batchCodeTable)
+        batchCodeTable = batchCodeTable[
+            [
+                'Adverse Reaction Reports',
+                'Deaths',
+                'Disabilities',
+                'Life Threatening Illnesses',
+                'Company',
+                'Countries',
+                'Severe reports',
+                'Lethality'
+            ]]
+        return batchCodeTable.sort_values(by = 'Severe reports', ascending = False)
+
+    def _getBatchCodeTableByCountry(self, country):
+        if country in self.countryBatchCodeTable.index:
+            return self.countryBatchCodeTable.loc[country]
+        else:
+            return self._getEmptyBatchCodeTable()
+
+    def _getEmptyBatchCodeTable(self):
+        return self.countryBatchCodeTable[0:0].droplevel(0)