Small fixes

.gitignore

@@ -1 +1,3 @@
-workbench.log
+workbench.log
+.mypy_cache
+__pycache__

src/app.py

@@ -10,9 +10,9 @@
 from contextlib import asynccontextmanager
 
 from src.database.database import BaseMongoDbWrapper
-import src.routers.employee_router as employee_router
-import src.routers.unit_router as unit_router
-import src.routers.workbench_router as workbench_router
+from src.routers import employee_router
+from src.routers import unit_router
+from src.routers import workbench_router
 from src._logging import HANDLERS
 from src.feecc_workbench.Messenger import MessageLevels, message_generator, messenger
 from src.database.models import GenericResponse
@@ -32,17 +32,17 @@ async def lifespan(app: FastAPI):
 
     yield
 
-    await WorkBench().shutdown()
+    await WorkBench.shutdown()
     BaseMongoDbWrapper.close_connection()
 
 
 # create app
 app = FastAPI(title="Feecc Workbench daemon", lifespan=lifespan)
 
 # include routers
-app.include_router(employee_router.router)
-app.include_router(unit_router.router)
-app.include_router(workbench_router.router)
+app.include_router(employee_router)
+app.include_router(unit_router)
+app.include_router(workbench_router)
 
 # set up CORS
 app.add_middleware(

src/database/database.py

@@ -26,6 +26,9 @@ def close_connection(self) -> None:
         self._client.close()
         logger.info("MongoDB connection closed")
 
+    def create_index(self, collection: str, index_name: str) -> None:
+        self._database[collection].create_index(index_name)
+
     def insert(self, collection: str, entity: dict[str, Any]) -> None:
         """Inserts the entity in the specified collection."""
         self._database[collection].insert_one(entity)

src/database/models.py

@@ -111,7 +111,6 @@ class HidEvent(BaseModel):
 
 class ProductionSchemaStage(BaseModel):
     name: str
-    stage_id: str
     type: str | None = None  # noqa: A003
     description: str | None = None
     equipment: list[str] | None = None
@@ -121,14 +120,13 @@ class ProductionSchemaStage(BaseModel):
 
 class ProductionSchema(BaseModel):
     schema_id: str = Field(default_factory=lambda: uuid4().hex)
-    unit_name: str
-    unit_short_name: str | None = None
+    schema_name: str
+    schema_print_name: str | None = None
     production_stages: list[ProductionSchemaStage]
-    required_components_schema_ids: list[str] | None = None
+    components_schema_ids: list[str] | None = None
     parent_schema_id: str | None = None
     schema_type: str | None = None
     erp_metadata: dict[str, str] | None = None
-    allowed_positions: list[str] | None = None
 
     @property
     def is_composite(self) -> bool:
@@ -140,9 +138,9 @@ def is_a_component(self) -> bool:
 
     @property
     def print_name(self) -> str:
-        if self.unit_short_name is None:
-            return self.unit_name
-        return self.unit_short_name
+        if self.schema_print_name is None:
+            return self.schema_name
+        return self.schema_print_name
 
     def is_allowed(self, position: str) -> bool:
         if not self.allowed_positions:

src/feecc_workbench/WorkBench.py

@@ -31,7 +31,7 @@
 STATE_SWITCH_EVENT = asyncio.Event()
 
 
-class WorkBench:
+class _WorkBench:
     """
     Work bench is a union of an Employee, working at it and Camera attached.
     It provides highly abstract interface for interaction with them
@@ -379,3 +379,6 @@ async def shutdown(self) -> None:
         message = "Workbench shutdown sequence complete"
         logger.info(message)
         messenger.success(translation("FinishServer"))
+
+
+WorkBench = _WorkBench()

src/feecc_workbench/_label_generation.py

@@ -116,18 +116,18 @@ def __init__(self, unit_code: str) -> None:
         self.barcode: barcode.EAN13 = barcode.get("ean13", self.unit_code, writer=ImageWriter())
         self.basename: str = f"output/barcode/{self.barcode.get_fullcode()}_barcode"
         self.filename: str = f"{self.basename}.png"
-        self.save_barcode(self.barcode)
-
-    def save_barcode(self, ean_code: barcode.EAN13) -> str:
-        """Method that saves the barcode image"""
-        dir_ = pathlib.Path(self.filename).parent
-        if not dir_.is_dir():
-            dir_.mkdir(parents=True)
-        barcode_path = str(
-            ean_code.save(self.basename, {"module_height": 12, "text_distance": 3, "font_size": 8, "quiet_zone": 1})
-        )
-        with Image.open(barcode_path) as img:
-            img = _resize_to_paper_aspect_ratio(img)
-            img.save(barcode_path)
-
-        return barcode_path
+
+
+def save_barcode(barcode: Barcode) -> str:
+    """Method that saves the barcode image"""
+    dir_ = pathlib.Path(barcode.filename).parent
+    if not dir_.is_dir():
+        dir_.mkdir(parents=True)
+    barcode_path = str(
+        barcode.ean_code.save(barcode.basename, {"module_height": 12, "text_distance": 3, "font_size": 8, "quiet_zone": 1})
+    )
+    with Image.open(barcode_path) as img:
+        img = _resize_to_paper_aspect_ratio(img)
+        img.save(barcode_path)
+
+    return barcode_path

src/routers/__init__.py

@@ -0,0 +1,3 @@
+from .employee_router import router as employee_router
+from .unit_router import router as unit_router
+from .workbench_router import router as workbench_router

src/routers/employee_router.py

@@ -8,7 +8,7 @@
 from src.feecc_workbench.exceptions import StateForbiddenError
 from src.feecc_workbench.WorkBench import WorkBench
 
-WORKBENCH = WorkBench()
+WORKBENCH = WorkBench
 
 router = APIRouter(
     prefix="/employee",

src/routers/unit_router.py

@@ -9,7 +9,7 @@
 from src.unit.unit_utils import Unit
 from src.feecc_workbench.WorkBench import WorkBench
 
-WORKBENCH = WorkBench()
+WORKBENCH = WorkBench
 
 router = APIRouter(
     prefix="/unit",

src/routers/workbench_router.py

@@ -19,7 +19,7 @@
 from src.feecc_workbench.WorkBench import STATE_SWITCH_EVENT, WorkBench
 from src.config import CONFIG
 
-WORKBENCH = WorkBench()
+WORKBENCH = WorkBench
 
 router = APIRouter(
     prefix="/workbench",
@@ -152,7 +152,7 @@ def get_schema_list_entry(schema: mdl.ProductionSchema) -> mdl.SchemaListEntry:
         handled_schemas.add(schema.schema_id)
         return mdl.SchemaListEntry(
             schema_id=schema.schema_id,
-            schema_name=schema.unit_name,
+            schema_name=schema.schema_name,
             included_schemas=included_schemas,
         )
 

src/unit/UnitManager.py

@@ -182,8 +182,8 @@ def assign_component(self, component: Unit) -> None:
         component.featured_in_int_id = self._get_cur_unit.internal_id
         logger.info(f"Component {component.model_name} has been assigned to a composite Unit {self.model_name}")
         messenger.success(
-            f"{translation("Component")} \
-{component.model_name} {translation("AssignedToUnit")} \
+            f"{translation('Component')} \
+{component.model_name} {translation('AssignedToUnit')} \
 {self.model_name}"
         )
 

src/unit/unit_utils.py

@@ -2,17 +2,22 @@
 
 import enum
 import datetime as dt
+
 from pydantic import BaseModel
 from uuid import uuid4
+from typing import TYPE_CHECKING
 
 from src.prod_stage.ProductionStage import ProductionStage
-from src.feecc_workbench._label_generation import Barcode
+from src.feecc_workbench._label_generation import Barcode, save_barcode
 from src.employee.Employee import Employee
 from src.prod_schema.prod_schema_wrapper import ProdSchemaWrapper
-from src.unit.unit_wrapper import UnitWrapper
 from src.database.models import ProductionSchema
 
 
+if TYPE_CHECKING:
+    from src.unit.unit_wrapper import UnitWrapper
+
+
 def biography_factory(schema_id: str, parent_unit_uuid: str) -> list[ProductionStage]:
     operation_stages = []
     production_schema = ProdSchemaWrapper.get_schema_by_id(schema_id)
@@ -72,12 +77,14 @@ def _get_unit_dict_data(unit: Unit) -> dict[str, str | bool | None | list[str] |
 
 
 class Unit(BaseModel):
+    class Config:
+        arbitrary_types_allowed=True
+
     status: UnitStatus
     schema_id: str  # The id of the schema used in production
     uuid: str = uuid4().hex
     operation_name: str  # The name of the operation (simple, complex etc)
     barcode: Barcode = Barcode(str(int(uuid, 16))[:12])
-    internal_id: str = str(barcode.barcode.get_fullcode())
     schema: ProductionSchema | None = None  # Used for initialization
     components_units: list[Unit] | None = None
     certificate_ipfs_cid: str | None = None
@@ -108,6 +115,9 @@ def model_post_init(self, __context: enum.Any) -> None:
 
         self._component_slots: dict[str, Unit | None] = slots
 
+        self.internal_id: str = str(self.barcode.barcode.get_fullcode())
+        save_barcode(self.barcode)
+
         if not self.operation_stages:
             self.operation_stages = biography_factory(self.schema_id, self.uuid)