adding CountriesColumnAdder
This commit is contained in:
8
src/CountriesColumnAdder.py
Normal file
8
src/CountriesColumnAdder.py
Normal file
@@ -0,0 +1,8 @@
|
||||
class CountriesColumnAdder:
|
||||
|
||||
def addCountriesColumn(self, countriesByBatchcodeTable):
|
||||
countriesByBatchcodeTable['Countries'] = countriesByBatchcodeTable.apply(self._getCountriesHavingEvents, axis='columns')
|
||||
return countriesByBatchcodeTable
|
||||
|
||||
def _getCountriesHavingEvents(self, eventCountByCountry):
|
||||
return set(eventCountByCountry[eventCountByCountry > 0].index)
|
||||
30
src/CountriesColumnAdderTest.py
Normal file
30
src/CountriesColumnAdderTest.py
Normal file
@@ -0,0 +1,30 @@
|
||||
import unittest
|
||||
from TestHelper import TestHelper
|
||||
from pandas.testing import assert_frame_equal
|
||||
import pandas as pd
|
||||
from src.CountriesColumnAdder import CountriesColumnAdder
|
||||
|
||||
class CountriesColumnAdderTest(unittest.TestCase):
|
||||
|
||||
def test_addCountriesColumn(self):
|
||||
# Given
|
||||
countriesByBatchcodeTable = TestHelper.createDataFrame(
|
||||
columns = ['United States', 'Germany', 'Italy'],
|
||||
data = [ [20, 0, 3]],
|
||||
index = pd.Index(
|
||||
name = 'Batchcode',
|
||||
data = ['FE6208']))
|
||||
|
||||
# When
|
||||
countriesByBatchcodeTableWithCountriesColumn = CountriesColumnAdder().addCountriesColumn(countriesByBatchcodeTable)
|
||||
|
||||
# Then
|
||||
assert_frame_equal(
|
||||
countriesByBatchcodeTableWithCountriesColumn,
|
||||
TestHelper.createDataFrame(
|
||||
columns = ['United States', 'Germany', 'Italy', 'Countries'],
|
||||
data = [ [20, 0, 3, {'United States', 'Italy'}]],
|
||||
index = pd.Index(
|
||||
name = 'Batchcode',
|
||||
data = ['FE6208'])))
|
||||
|
||||
@@ -35,6 +35,8 @@
|
||||
" exploration.drop(index=indexName, inplace=True)\n",
|
||||
" exploration.index.rename(indexName, inplace=True)\n",
|
||||
" exploration.drop(columns='Totals', inplace=True)\n",
|
||||
" for column in exploration.columns:\n",
|
||||
" exploration[column] = exploration[column].astype('int64')\n",
|
||||
" return exploration"
|
||||
]
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user