From 92ea72249cd8ccc68d7786e453382b1de091fcbb Mon Sep 17 00:00:00 2001
From: TrisNol <tristan.nolde@yahoo.de>
Date: Sun, 17 Sep 2023 19:43:22 +0200
Subject: [PATCH] refactor(data-extraction): Bind company type to an enum

---
 .../models/company.py                         | 24 ++++++++++++++++++-
 .../models/temp.md                            |  0
 .../unternehmensregister/transform.py         | 24 +++++++++++--------
 3 files changed, 37 insertions(+), 11 deletions(-)
 delete mode 100644 src/aki_prj23_transparenzregister/models/temp.md

diff --git a/src/aki_prj23_transparenzregister/models/company.py b/src/aki_prj23_transparenzregister/models/company.py
index 66f671a..67806a9 100644
--- a/src/aki_prj23_transparenzregister/models/company.py
+++ b/src/aki_prj23_transparenzregister/models/company.py
@@ -31,6 +31,28 @@ class RelationshipRoleEnum(str, MultiValueEnum):
     HAUPTNIEDERLASSUNG = "Hauptniederlassung"
 
 
+class CompanyTypeEnum(str, MultiValueEnum):
+    """Type of Company."""
+
+    GMBH = "Gesellschaft mit beschränkter Haftung"
+    SE = "Europäische Aktiengesellschaft (SE)"
+    KG = "Kommanditgesellschaft"
+    EINZELKAUFMANN = (
+        "Einzelkaufmann",
+        "Einzelkauffrau",
+        "Einzelkaufmann / Einzelkauffrau",
+    )
+    EG = "eingetragene Genossenschaft"
+    AG = "Aktiengesellschaft"
+    PARTNERSCHAFTSGESELLSCHAFT = "Partnerschaftsgesellschaft"
+    PARTNERGESELLSCHAFT = "Partnergesellschaft"
+    PARTNERSCHAFT = "Partnerschaft"
+    KGaA = "Kommanditgesellschaft auf Aktien"
+    OHG = "Offene Handelsgesellschaft"
+    AUSLAENDISCHE_RECHTSFORM = "Rechtsform ausländischen Rechts HRB"
+    JURISTISCHE_PERSON = "HRA Juristische Person"
+
+
 @dataclass
 class CompanyID:
     """_summary_."""
@@ -153,7 +175,7 @@ class Company:
     last_update: str
     relationships: list[CompanyRelationship]
     # yearly_results: Optional[list[FinancialResults]]
-    company_type: str | None = None  # TODO define Enum
+    company_type: CompanyTypeEnum | None = None
     capital: Capital | None | None = None
     business_purpose: str | None = None
     founding_date: str | None = None
diff --git a/src/aki_prj23_transparenzregister/models/temp.md b/src/aki_prj23_transparenzregister/models/temp.md
deleted file mode 100644
index e69de29..0000000
diff --git a/src/aki_prj23_transparenzregister/utils/data_extraction/unternehmensregister/transform.py b/src/aki_prj23_transparenzregister/utils/data_extraction/unternehmensregister/transform.py
index 8f5a51c..a95667f 100644
--- a/src/aki_prj23_transparenzregister/utils/data_extraction/unternehmensregister/transform.py
+++ b/src/aki_prj23_transparenzregister/utils/data_extraction/unternehmensregister/transform.py
@@ -16,6 +16,7 @@ from aki_prj23_transparenzregister.models.company import (
     CompanyRelationship,
     CompanyRelationshipEnum,
     CompanyToCompanyRelationship,
+    CompanyTypeEnum,
     Location,
     PersonName,
     PersonToCompanyRelationship,
@@ -205,7 +206,7 @@ def name_from_beteiligung(data: dict) -> str:
     ]["Organisation"]["Bezeichnung"]["Bezeichnung_Aktuell"]
 
 
-def map_rechtsform(company_name: str, data: dict) -> str | None:
+def map_rechtsform(company_name: str, data: dict) -> CompanyTypeEnum | None:
     """Extracts the company type from a given Unternehmensregister export.
 
     Args:
@@ -213,23 +214,25 @@ def map_rechtsform(company_name: str, data: dict) -> str | None:
         data (dict): Data export
 
     Returns:
-        str | None: Company type if found
+        CompanyTypeEnum | None: Company type if found
     """
     try:
-        return data["XJustiz_Daten"]["Fachdaten_Register"]["Basisdaten_Register"][
-            "Rechtstraeger"
-        ]["Rechtsform"]["content"]
-    except Exception:
+        return CompanyTypeEnum(
+            data["XJustiz_Daten"]["Fachdaten_Register"]["Basisdaten_Register"][
+                "Rechtstraeger"
+            ]["Rechtsform"]["content"]
+        )
+    except KeyError:
         if (
             company_name.endswith("GmbH")
             or company_name.endswith("UG")
             or company_name.endswith("UG (haftungsbeschränkt)")
         ):
-            return "Gesellschaft mit beschränkter Haftung"
+            return CompanyTypeEnum("Gesellschaft mit beschränkter Haftung")
         if company_name.endswith("SE"):
-            return "Europäische Aktiengesellschaft (SE)"
+            return CompanyTypeEnum("Europäische Aktiengesellschaft (SE)")
         if company_name.endswith("KG"):
-            return "Kommanditgesellschaft"
+            return CompanyTypeEnum("Kommanditgesellschaft")
         return None
 
 
@@ -487,5 +490,6 @@ if __name__ == "__main__":
                         dataclasses.asdict(company), export_file, ensure_ascii=False
                     )
             except Exception as e:
-                logger.error(f"Error in processing {path}\n{e}")
+                logger.error(f"Error in processing {path}")
+                logger.error(e)
                 break