code cleanup

dam_helper/local_storage/entities.py

@@ -96,3 +96,10 @@ class ContractorDetails(SavedDetails):
         if self.contractor_name is not None:
             answers[Questions.COMPANY_NAME] = self.contractor_name
         return answers
+
+
+detail_classes = {
+    DetailsType.LOCATION_DETAILS: LocationDetails,
+    DetailsType.PERSONAL_DETAILS: PersonalDetails,
+    DetailsType.CONTRACTOR_DETAILS: ContractorDetails
+}

dam_helper/local_storage/ls_manager.py

@@ -1,6 +1,7 @@
 from streamlit_local_storage import LocalStorage
 
-from dam_helper.local_storage.entities import SavedDetails, PersonalDetails, LocationDetails, ContractorDetails
+from dam_helper.local_storage.entities import SavedDetails, PersonalDetails, LocationDetails, ContractorDetails, \
+    detail_classes
 from dam_helper.enums.enums import DetailsType
 
 
@@ -12,34 +13,17 @@ class LocalStorageManager:
         for detail_type in DetailsType.values():
             detail = self.ls.getItem(detail_type.name).get(key)
             if detail:
-                if detail_type.value == DetailsType.PERSONAL_DETAILS.value:
-                    return PersonalDetails(**detail)
-                elif detail_type.value == DetailsType.LOCATION_DETAILS.value:
-                    return LocationDetails(**detail)
-                elif detail_type.value == DetailsType.CONTRACTOR_DETAILS.value:
-                    return ContractorDetails(**detail)
+                cls = detail_classes[detail_type]
+                cls(**detail)
                 return None
         return None
 
     def save_details(self, details: SavedDetails):
-        if isinstance(details, PersonalDetails):
-            type_ = DetailsType.PERSONAL_DETAILS
-        elif isinstance(details, LocationDetails):
-            type_ = DetailsType.LOCATION_DETAILS
-        elif isinstance(details, ContractorDetails):
-            type_ = DetailsType.CONTRACTOR_DETAILS
-        else:
-            raise ValueError("Unexpected type: {}", type(details))
-        existing_data = (self.ls.getItem(type_.name) or {})
+        existing_data = (self.ls.getItem(details.type_.name) or {})
         existing_data[details.key] = details.to_json()
-        self.ls.setItem(type_.name, existing_data, type_.name)
+        self.ls.setItem(details.type_.name, existing_data, details.type_.name)
 
     def get_details(self, type_: DetailsType) -> list[SavedDetails]:
-        if type_.value == DetailsType.PERSONAL_DETAILS.value:
-            cls = PersonalDetails
-        elif type_.value == DetailsType.LOCATION_DETAILS.value:
-            cls = LocationDetails
-        else:
-            cls = ContractorDetails
+        cls = detail_classes[type_]
         details = self.ls.getItem(type_.name) or {}
         return [cls(**v) for k,v in details.items()]

dam_helper/ui_manager/ui_manager.py

@@ -8,7 +8,8 @@ from dam_helper.email_manager.email_manager import EmailManager
 from dam_helper.enums.enums import Steps, DetailsType, Questions as Q
 from dam_helper.form.form import work_categories, build_form_data_from_answers, write_pdf_form
 from dam_helper.llm_manager.llm_parser import LlmParser
-from dam_helper.local_storage.entities import PersonalDetails, LocationDetails, ContractorDetails, SavedDetails
+from dam_helper.local_storage.entities import PersonalDetails, LocationDetails, ContractorDetails, SavedDetails, \
+    detail_classes
 from dam_helper.local_storage.ls_manager import LocalStorageManager
 from dam_helper.prompts.prompts_manager import PromptsManager
 from dam_helper.repository import Model, ModelRoles
@@ -228,12 +229,7 @@ class UIManager:
 
     @st.dialog("Add new")
     def add_new_detail_dialog(self, type_: DetailsType):
-        if type_.name == DetailsType.CONTRACTOR_DETAILS.name:
-            new_item = ContractorDetails()
-        elif type_.name == DetailsType.PERSONAL_DETAILS.name:
-            new_item = PersonalDetails()
-        else:
-            new_item = LocationDetails()
+        new_item = detail_classes[type_]()
         with st.form("new item", border=False):
             fields_labels = new_item.widget_labels()
             for k, v in fields_labels.items():