Add types and ignores to make mypy happy

src/cyhy_db/db.py

@@ -1,10 +1,10 @@
 # Third-Party Libraries
-from beanie import init_beanie
-from motor.motor_asyncio import AsyncIOMotorClient
+from beanie import Document, View, init_beanie
+from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
 
 from .models import *
 
-ALL_MODELS = [
+ALL_MODELS: list[type[Document] | type[View] | str] = [
     CVE,
     HostDoc,
     HostScanDoc,
@@ -22,10 +22,10 @@
 ]
 
 
-async def initialize_db(db_uri: str, db_name: str) -> None:
+async def initialize_db(db_uri: str, db_name: str) -> AsyncIOMotorDatabase:
     try:
-        client = AsyncIOMotorClient(db_uri)
-        db = client[db_name]
+        client: AsyncIOMotorClient = AsyncIOMotorClient(db_uri)
+        db: AsyncIOMotorDatabase = client[db_name]
         await init_beanie(database=db, document_models=ALL_MODELS)
         return db
     except Exception as e:

src/cyhy_db/models/cve.py

@@ -12,7 +12,8 @@ class CVE(Document):
     # Validate on assignment so ip_int is recalculated as ip is set
     model_config = ConfigDict(extra="forbid", validate_assignment=True)
 
-    id: str = Indexed(primary_field=True)  # CVE ID
+    # CVE ID as a string
+    id: str = Indexed(primary_field=True)  # type: ignore[assignment]
     cvss_score: float = Field(ge=0.0, le=10.0)
     cvss_version: CVSSVersion = Field(default=CVSSVersion.V3_1)
     severity: int = Field(ge=1, le=4, default=1)

src/cyhy_db/models/host_doc.py

@@ -21,7 +21,8 @@ class State(BaseModel):
 class HostDoc(Document):
     model_config = ConfigDict(extra="forbid")
 
-    id: int = Field()  # IP address as an integer
+    # IP address as an integer
+    id: int = Field(default_factory=int)  # type: ignore[assignment]
     ip: IPv4Address = Field(...)
     owner: str = Field(...)
     last_change: datetime = Field(default_factory=utcnow)

src/cyhy_db/models/kev_doc.py

@@ -1,12 +1,12 @@
 # Third-Party Libraries
 from beanie import Document
-from pydantic import ConfigDict
+from pydantic import ConfigDict, Field
 
 
 class KEVDoc(Document):
     model_config = ConfigDict(extra="forbid")
 
-    id: str  # CVE
+    id: str = Field(default_factory=str)  # type: ignore[assignment]
     known_ransomware: bool
 
     class Settings:

src/cyhy_db/models/place_doc.py

@@ -2,14 +2,15 @@
 from typing import Optional
 
 # Third-Party Libraries
-from beanie import Document
+from beanie import Document, PydanticObjectId
 from pydantic import ConfigDict, Field
 
 
 class PlaceDoc(Document):
     model_config = ConfigDict(extra="forbid")
 
-    id: int  # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html
+    # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html
+    id: int = Field(default_factory=int)  # type: ignore[assignment]
     name: str
     clazz: str = Field(alias="class")  # 'class' is a reserved keyword in Python
     state: str

src/cyhy_db/models/request_doc.py

@@ -5,7 +5,15 @@
 from typing import List, Optional
 
 # Third-Party Libraries
-from beanie import Document, Insert, Link, Replace, ValidateOnSave, before_event
+from beanie import (
+    Document,
+    Insert,
+    Link,
+    PydanticObjectId,
+    Replace,
+    ValidateOnSave,
+    before_event,
+)
 from pydantic import BaseModel, ConfigDict, EmailStr, Field, field_validator
 
 from ..utils import utcnow
@@ -81,7 +89,7 @@ def validate_start(cls, v):
 class RequestDoc(Document):
     model_config = ConfigDict(extra="forbid")
 
-    id: str = Field(default=BOGUS_ID)
+    id: str = Field(default=BOGUS_ID)  # type: ignore[assignment]
     agency: Agency
     children: List[Link["RequestDoc"]] = Field(default=[])
     enrolled: datetime = Field(default_factory=utcnow)

src/cyhy_db/models/system_control_doc.py

@@ -33,7 +33,7 @@ async def wait_for_completion(cls, document_id, timeout: Optional[int] = None):
         start_time = utcnow()
         while True:
             doc = await cls.get(document_id)
-            if doc.completed:
+            if doc and doc.completed:
                 return True
             if timeout and (utcnow() - start_time).total_seconds() > timeout:
                 return False