From 971f01935a811138bc40a25bf9e3b989a8330ebb Mon Sep 17 00:00:00 2001
From: frankknoll <Knoll_Frank@web.de>
Date: Mon, 6 Feb 2023 15:58:32 +0100
Subject: [PATCH] refactoring

---
 src/CountryColumnAdder.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/src/CountryColumnAdder.py b/src/CountryColumnAdder.py
index f7d14ccd4ab..9e703ad696c 100644
--- a/src/CountryColumnAdder.py
+++ b/src/CountryColumnAdder.py
@@ -4,17 +4,18 @@ class CountryColumnAdder:
     
     @staticmethod
     def addCountryColumn(dataFrame):
-        dataFrame['COUNTRY'] = CountryColumnAdder.getCountryColumn(dataFrame)
-        return dataFrame.astype({'COUNTRY': "string"})
+        dataFrame['COUNTRY'] = CountryColumnAdder._getCountryColumn(dataFrame)
+        return dataFrame
 
     @staticmethod
-    def getCountryColumn(dataFrame):
-        return dataFrame.apply(
-            lambda row:
-                CountryColumnAdder._getCountryNameOfSplttypeOrDefault(
-                 splttype = row['SPLTTYPE'],
-                 default = 'Unknown Country'),
-            axis = 'columns')
+    def _getCountryColumn(dataFrame):
+        return (dataFrame['SPLTTYPE']
+                .apply(
+                    lambda splttype:
+                        CountryColumnAdder._getCountryNameOfSplttypeOrDefault(
+                            splttype = splttype,
+                            default = 'Unknown Country'))
+                .astype("string"))
 
     @staticmethod
     def _getCountryNameOfSplttypeOrDefault(splttype, default):