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.drop(index=indexName, inplace=True)\n",
|
||||||
" exploration.index.rename(indexName, inplace=True)\n",
|
" exploration.index.rename(indexName, inplace=True)\n",
|
||||||
" exploration.drop(columns='Totals', 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"
|
" return exploration"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
|||||||
Reference in New Issue
Block a user