refactor(data-extraction): Bind company type to an enum

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

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"][
+        return CompanyTypeEnum(
-            "Rechtstraeger"
+            data["XJustiz_Daten"]["Fachdaten_Register"]["Basisdaten_Register"][
-        ]["Rechtsform"]["content"]
+                "Rechtstraeger"
-    except Exception:
+            ]["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