Add type annotations to ORM methods

datafiles/managers.py

@@ -5,6 +5,7 @@
 import os
 from pathlib import Path
 from typing import Any, Dict, Optional
+from typing_extensions import Protocol
 
 import log
 from cached_property import cached_property
@@ -18,35 +19,6 @@
 Missing = dataclasses._MISSING_TYPE
 
 
-class Manager:
-    def __init__(self, cls):
-        self.model = cls
-
-    def all(self):
-        raise NotImplementedError
-
-    def get_or_none(self, *args, **kwargs):
-        original_manual = self.model.Meta.datafile_manual
-
-        self.model.Meta.datafile_manual = True
-        instance = self.model(*args, **kwargs)
-        self.model.Meta.datafile_manual = original_manual
-
-        if instance.datafile.exists:
-            instance.datafile._manual = original_manual
-            return instance
-
-        return None
-
-    def get_or_create(self, *args, **kwargs):
-        instance = self.model(*args, **kwargs)
-
-        if not instance.datafile.exists:
-            instance.datafile.save()
-
-        return instance
-
-
 class Datafile:
     def __init__(
         self,
@@ -368,3 +340,37 @@ def _write(self, text: str):
         log.debug('=' * len(message) + '\n\n' + (text or '<nothing>\n'))
         self.path.parent.mkdir(parents=True, exist_ok=True)
         self.path.write_text(text)
+
+
+class HasDatafile(Protocol):
+
+    datafile: Datafile
+
+
+class Manager:
+    def __init__(self, cls):
+        self.model = cls
+
+    def all(self):
+        raise NotImplementedError
+
+    def get_or_none(self, *args, **kwargs) -> Optional[HasDatafile]:
+        original_manual = self.model.Meta.datafile_manual
+
+        self.model.Meta.datafile_manual = True
+        instance = self.model(*args, **kwargs)
+        self.model.Meta.datafile_manual = original_manual
+
+        if instance.datafile.exists:
+            instance.datafile._manual = original_manual
+            return instance
+
+        return None
+
+    def get_or_create(self, *args, **kwargs) -> HasDatafile:
+        instance = self.model(*args, **kwargs)
+
+        if not instance.datafile.exists:
+            instance.datafile.save()
+
+        return instance

pyproject.toml

@@ -49,6 +49,9 @@ minilog = "^1.2.3"
 classproperties = "~0.1.3"
 cached_property = "^1.5"
 
+# Typing
+typing-extensions = "^3.7"
+
 [tool.poetry.dev-dependencies]
 
 # Formatters