2 files changed, 14 insertions, 13 deletions
diff --git a/python/fatcat_tools/cleanups/common.py b/python/fatcat_tools/cleanups/common.py
index 26ca7bd6..7ebfc8a0 100644
--- a/python/fatcat_tools/cleanups/common.py
+++ b/python/fatcat_tools/cleanups/common.py
@@ -2,6 +2,7 @@ import copy
 import json
 import subprocess
 from collections import Counter
+from typing import Any, Dict, List
 
 from fatcat_openapi_client import ApiClient, Editgroup
 
@@ -27,7 +28,7 @@ class EntityCleaner:
     This class is pretty similar to EntityImporter, but isn't subclassed.
     """
 
-    def __init__(self, api, entity_type, **kwargs):
+    def __init__(self, api: ApiClient, entity_type: Any, **kwargs) -> None:
 
         eg_extra = kwargs.get("editgroup_extra", dict())
         eg_extra["git_rev"] = eg_extra.get(
@@ -49,14 +50,14 @@ class EntityCleaner:
         if self.dry_run_mode:
             print("Running in dry-run mode!")
 
-    def reset(self):
+    def reset(self) -> None:
         self.counts = Counter({"lines": 0, "cleaned": 0, "updated": 0})
         self._edit_count = 0
         self._editgroup_id = None
-        self._entity_queue = []
-        self._idents_inflight = []
+        self._entity_queue: List[Any] = []
+        self._idents_inflight: List[str] = []
 
-    def push_record(self, record):
+    def push_record(self, record: Dict[str, Any]) -> None:
         """
         Intended to be called by "pusher" class (which could be pulling from
         JSON file, Kafka, whatever).
@@ -106,14 +107,14 @@ class EntityCleaner:
             self._idents_inflight = []
         return
 
-    def clean_entity(self, entity):
+    def clean_entity(self, entity: Any) -> Any:
         """
         Mutates entity in-place and returns it
         """
         # implementations should fill this in
         raise NotImplementedError
 
-    def try_update(self, entity):
+    def try_update(self, entity: Any) -> int:
         """
         Returns edit count (number of entities updated).
 
@@ -123,7 +124,7 @@ class EntityCleaner:
         # implementations should fill this in
         raise NotImplementedError
 
-    def finish(self):
+    def finish(self) -> Counter:
         if self._edit_count > 0:
             self.api.accept_editgroup(self._editgroup_id)
             self._editgroup_id = None
@@ -132,7 +133,7 @@ class EntityCleaner:
 
         return self.counts
 
-    def get_editgroup_id(self):
+    def get_editgroup_id(self) -> str:
 
         if not self._editgroup_id:
             eg = self.api.create_editgroup(
diff --git a/python/fatcat_tools/cleanups/files.py b/python/fatcat_tools/cleanups/files.py
index d378a91f..309924e3 100644
--- a/python/fatcat_tools/cleanups/files.py
+++ b/python/fatcat_tools/cleanups/files.py
@@ -1,4 +1,4 @@
-from fatcat_openapi_client.models import FileEntity
+from fatcat_openapi_client import ApiClient, FileEntity
 from fatcat_openapi_client.rest import ApiException
 
 from .common import EntityCleaner
@@ -9,7 +9,7 @@ class FileCleaner(EntityCleaner):
     File fixups!
     """
 
-    def __init__(self, api, **kwargs):
+    def __init__(self, api: ApiClient, **kwargs) -> None:
 
         eg_desc = (
             kwargs.pop("editgroup_description", None)
@@ -25,7 +25,7 @@ class FileCleaner(EntityCleaner):
             **kwargs
         )
 
-    def clean_entity(self, entity):
+    def clean_entity(self, entity: FileEntity) -> FileEntity:
         """
         TODO: mimetype is bogus like (???) => clean mimetype
         """
@@ -54,7 +54,7 @@ class FileCleaner(EntityCleaner):
 
         return entity
 
-    def try_update(self, entity):
+    def try_update(self, entity: FileEntity) -> int:
 
         try:
             existing = self.api.get_file(entity.ident)