Mypy backlog cleanup (#601)

* Errors under 90

* More mypy

* Fixing the rest of mypy things

examples/parlai_chat_task_demo/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

examples/parlai_chat_task_demo/demo_worlds.py

@@ -3,9 +3,9 @@
 # Copyright (c) Facebook, Inc. and its affiliates.
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-from parlai.crowdsourcing.utils.worlds import CrowdOnboardWorld, CrowdTaskWorld
-from parlai.core.worlds import validate
-from joblib import Parallel, delayed
+from parlai.crowdsourcing.utils.worlds import CrowdOnboardWorld, CrowdTaskWorld  # type: ignore
+from parlai.core.worlds import validate  # type: ignore
+from joblib import Parallel, delayed  # type: ignore
 
 
 class MultiAgentDialogOnboardWorld(CrowdOnboardWorld):

examples/simple_static_task/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

examples/simple_static_task/examine_results.py

@@ -8,6 +8,7 @@
 from mephisto.abstractions.databases.local_database import LocalMephistoDB
 from mephisto.tools.examine_utils import run_examine_or_review, print_results
 from mephisto.data_model.worker import Worker
+from mephisto.data_model.unit import Unit
 
 db = None
 

examples/static_react_task/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

mephisto/abstractions/architects/channels/websocket_channel.py

@@ -9,7 +9,7 @@
 from mephisto.abstractions.channel import Channel, STATUS_CHECK_TIME
 
 import errno
-import websocket
+import websocket  # type: ignore
 import threading
 import json
 import time

mephisto/abstractions/architects/heroku_architect.py

@@ -11,7 +11,7 @@
 import netrc
 import os
 import platform
-import sh
+import sh  # type: ignore
 import shlex
 import shutil
 import subprocess

mephisto/abstractions/architects/local_architect.py

@@ -7,7 +7,7 @@
 import os
 import signal
 import subprocess
-import sh
+import sh  # type: ignore
 import shutil
 import shlex
 import time
@@ -161,7 +161,9 @@ def deploy(self) -> str:
                 preexec_fn=os.setpgrp,
                 env=dict(os.environ, PORT=f"{self.port}"),
             )
-        self.server_process_pid = self.server_process.pid
+        my_process = self.server_process
+        assert my_process is not None, "Cannot start without a process..."
+        self.server_process_pid = my_process.pid
         os.chdir(return_dir)
 
         time.sleep(1)

mephisto/abstractions/architects/router/build_router.py

@@ -6,7 +6,7 @@
 
 import mephisto.abstractions.architects.router as router_module
 import os
-import sh
+import sh  # type: ignore
 import shutil
 import shlex
 import subprocess

mephisto/abstractions/architects/router/flask/app.py

@@ -4,25 +4,25 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from gevent import monkey
+from gevent import monkey  # type: ignore
 
 monkey.patch_all()
 
 try:
-    from mephisto.abstractions.architects.router.flask.mephisto_flask_blueprint import (
+    from mephisto.abstractions.architects.router.flask.mephisto_flask_blueprint import (  # type: ignore
         MephistoRouter,
         mephisto_router,
     )
 except:
-    from mephisto_flask_blueprint import (
+    from mephisto_flask_blueprint import (  # type: ignore
         MephistoRouter,
         mephisto_router,
     )
-from geventwebsocket import WebSocketServer, Resource
-from werkzeug.debug import DebuggedApplication
+from geventwebsocket import WebSocketServer, Resource  # type: ignore
+from werkzeug.debug import DebuggedApplication  # type: ignore
 
 
-from flask import Flask
+from flask import Flask  # type: ignore
 import os
 
 port = int(os.environ.get("PORT", 3000))

mephisto/abstractions/architects/router/flask/mephisto_flask_blueprint.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from flask import (
+from flask import (  # type: ignore
     Flask,
     Blueprint,
     render_template,
@@ -13,7 +13,7 @@
     send_from_directory,
     jsonify,
 )
-from geventwebsocket import (
+from geventwebsocket import (  # type: ignore
     WebSocketServer,
     WebSocketApplication,
     Resource,
@@ -23,15 +23,15 @@
 import time
 import json
 import os
-from werkzeug.utils import secure_filename
+from werkzeug.utils import secure_filename  # type: ignore
 
 from threading import Event
 
 from typing import Dict, Tuple, List, Any, Optional, TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from geventwebsocket.handler import Client
-    from geventwebsocket.websocket import WebSocket
+    from geventwebsocket.handler import Client  # type: ignore
+    from geventwebsocket.websocket import WebSocket  # type: ignore
 
 # Constants
 
@@ -396,6 +396,7 @@ def make_provider_request(
 def handle_provider_request(request_type: str, data: Dict[str, Any]):
     """Wrapper for provider requests that handles extracting the result and timing out"""
     provider_data = data["provider_data"]
+    assert mephisto_router_app is not None, "Must initialize router before this call"
     packet = mephisto_router_app.get_default_provider_request_packet(
         request_type, provider_data
     )

mephisto/abstractions/blueprint.py

@@ -197,12 +197,11 @@ def launch_unit(self, unit: "Unit", agent: "Agent") -> None:
             AgentShutdownError,
         ) as e:
             # A returned Unit can be worked on again by someone else.
-            if (
-                unit.get_status() != AssignmentState.EXPIRED
-                and unit.get_assigned_agent().db_id == agent.db_id
-            ):
-                logger.debug(f"Clearing {agent} from {unit} due to {e}")
-                unit.clear_assigned_agent()
+            if unit.get_status() != AssignmentState.EXPIRED:
+                unit_agent = unit.get_assigned_agent()
+                if unit_agent is not None and unit_agent.db_id == agent.db_id:
+                    logger.debug(f"Clearing {agent} from {unit} due to {e}")
+                    unit.clear_assigned_agent()
             self.cleanup_unit(unit)
         except Exception as e:
             logger.exception(f"Unhandled exception in unit {unit}: {repr(e)}")
@@ -539,29 +538,32 @@ class BlueprintMixin(ABC):
     work, such that only the highest class needs to be called.
     """
 
-    @property
-    @abstractmethod
-    def ArgsMixin(self) -> Any:  # Should be a dataclass, to extend BlueprintArgs
-        pass
+    # @property
+    # @abstractmethod
+    # def ArgsMixin(self) -> Type[object]:  # Should be a dataclass, to extend BlueprintArgs
+    #     pass
 
-    @property
-    @abstractmethod
-    def SharedStateMixin(
-        self,
-    ) -> Any:  # Also should be a dataclass, to extend SharedTaskState
-        pass
+    # @property
+    # @abstractmethod
+    # def SharedStateMixin(
+    #     self,
+    # ) -> Type[object]:  # Also should be a dataclass, to extend SharedTaskState
+    #     pass
+    ArgsMixin: ClassVar[Type[object]]
+    SharedStateMixin: ClassVar[Type[object]]
 
     @staticmethod
-    def extract_unique_mixins(blueprint_class: ClassVar[Type["Blueprint"]]):
+    def extract_unique_mixins(blueprint_class: Type["Blueprint"]):
         """Return the unique mixin classes that are used in the given blueprint class"""
         mixin_subclasses = [
             clazz
             for clazz in blueprint_class.mro()
             if issubclass(clazz, BlueprintMixin)
         ]
+        target_class: Union[Type["Blueprint"], Type["BlueprintMixin"]] = blueprint_class
         # Remove magic created with `mixin_args_and_state`
-        while blueprint_class.__name__ == "MixedInBlueprint":
-            blueprint_class = mixin_subclasses.pop(0)
+        while target_class.__name__ == "MixedInBlueprint":
+            target_class = mixin_subclasses.pop(0)
         removed_locals = [
             clazz
             for clazz in mixin_subclasses
@@ -570,7 +572,7 @@ def extract_unique_mixins(blueprint_class: ClassVar[Type["Blueprint"]]):
         filtered_subclasses = set(
             clazz
             for clazz in removed_locals
-            if clazz != BlueprintMixin and clazz != blueprint_class
+            if clazz != BlueprintMixin and clazz != target_class
         )
         # we also want to make sure that we don't double-count extensions of mixins, so remove classes that other classes are subclasses of
         def is_subclassed(clazz):
@@ -607,7 +609,9 @@ def get_mixin_qualifications(
         raise NotImplementedError()
 
     @classmethod
-    def mixin_args_and_state(mixin_cls: "BlueprintMixin", target_cls: "Blueprint"):
+    def mixin_args_and_state(
+        mixin_cls: Type["BlueprintMixin"], target_cls: Type["Blueprint"]
+    ):
         """
         Magic utility decorator that can be used to inject mixin configurations
         (BlueprintArgs and SharedTaskState) without the user needing to define new
@@ -620,18 +624,18 @@ def mixin_args_and_state(mixin_cls: "BlueprintMixin", target_cls: "Blueprint"):
           class MyBlueprint(ABlueprintMixin, Blueprint):
               pass
         """
-
+        # Ignore typing on most of this, as mypy is not able to parse what's happening
         @dataclass
-        class MixedInArgsClass(mixin_cls.ArgsMixin, target_cls.ArgsClass):
+        class MixedInArgsClass(mixin_cls.ArgsMixin, target_cls.ArgsClass):  # type: ignore
             pass
 
         @dataclass
         class MixedInSharedStateClass(
-            mixin_cls.SharedStateMixin, target_cls.SharedStateClass
+            mixin_cls.SharedStateMixin, target_cls.SharedStateClass  # type: ignore
         ):
             pass
 
-        class MixedInBlueprint(target_cls):
+        class MixedInBlueprint(target_cls):  # type: ignore
             ArgsClass = MixedInArgsClass
             SharedStateClass = MixedInSharedStateClass
 

mephisto/abstractions/blueprints/abstract/static_task/static_blueprint.py

@@ -12,6 +12,7 @@
 from mephisto.abstractions.blueprints.mixins.onboarding_required import (
     OnboardingRequired,
     OnboardingSharedState,
+    OnboardingRequiredArgs,
 )
 from dataclasses import dataclass, field
 from omegaconf import MISSING, DictConfig
@@ -52,12 +53,12 @@
 
 
 @dataclass
-class SharedStaticTaskState(OnboardingRequired.SharedStateMixin, SharedTaskState):
+class SharedStaticTaskState(OnboardingSharedState, SharedTaskState):
     static_task_data: Iterable[Any] = field(default_factory=list)
 
 
 @dataclass
-class StaticBlueprintArgs(OnboardingRequired.ArgsMixin, BlueprintArgs):
+class StaticBlueprintArgs(OnboardingRequiredArgs, BlueprintArgs):
     _blueprint_type: str = BLUEPRINT_TYPE
     _group: str = field(
         default="StaticBlueprint",
@@ -131,7 +132,7 @@ def __init__(
         elif blue_args.get("data_json", None) is not None:
             json_file = os.path.expanduser(blue_args.data_json)
             with open(json_file, "r", encoding="utf-8-sig") as json_fp:
-                json_data = json.loads(json_fp)
+                json_data = json.load(json_fp)
             for jd in json_data:
                 self._initialization_data_dicts.append(jd)
         elif blue_args.get("data_jsonl", None) is not None:
@@ -149,10 +150,13 @@ def __init__(
             pass
 
     @classmethod
-    def assert_task_args(cls, args: DictConfig, shared_state: "SharedStaticTaskState"):
+    def assert_task_args(cls, args: DictConfig, shared_state: "SharedTaskState"):
         """Ensure that the data can be properly loaded"""
         super().assert_task_args(args, shared_state)
 
+        assert isinstance(
+            shared_state, SharedStaticTaskState
+        ), "Must use SharedStaticTaskState for static blueprints"
         blue_args = args.blueprint
         if blue_args.get("data_csv", None) is not None:
             csv_file = os.path.expanduser(blue_args.data_csv)
@@ -175,7 +179,7 @@ def assert_task_args(cls, args: DictConfig, shared_state: "SharedStaticTaskState
                 pass
             else:
                 assert (
-                    len(shared_state.static_task_data) > 0
+                    len([x for x in shared_state.static_task_data]) > 0
                 ), "Length of data dict provided was 0"
         else:
             raise AssertionError(

mephisto/abstractions/blueprints/mixins/onboarding_required.py

@@ -53,6 +53,7 @@ class OnboardingRequired(BlueprintMixin):
     Compositional class for blueprints that may have an onboarding step
     """
 
+    shared_state: "SharedTaskState"
     ArgsMixin = OnboardingRequiredArgs
     SharedStateMixin = OnboardingSharedState
 
@@ -100,6 +101,9 @@ def get_failed_qual(qual_name: str) -> str:
     def init_onboarding_config(
         self, task_run: "TaskRun", args: "DictConfig", shared_state: "SharedTaskState"
     ):
+        assert isinstance(
+            shared_state, OnboardingSharedState
+        ), f"Cannot init onboarding config with {shared_state}, need OnboardingSharedState"
         self.onboarding_qualification_name: Optional[str] = args.blueprint.get(
             "onboarding_qualification", None
         )
@@ -108,14 +112,16 @@ def init_onboarding_config(
         self.onboarding_qualification_id = None
         if not self.use_onboarding:
             return
+        onboarding_qualification_name = self.onboarding_qualification_name
+        assert onboarding_qualification_name is not None
 
         db = task_run.db
         self.onboarding_qualification_id = find_or_create_qualification(
             db,
-            self.onboarding_qualification_name,
+            onboarding_qualification_name,
         )
         self.onboarding_failed_name = self.get_failed_qual(
-            self.onboarding_qualification_name
+            onboarding_qualification_name
         )
         self.onboarding_failed_id = find_or_create_qualification(
             db, self.onboarding_failed_name
@@ -143,5 +149,9 @@ def validate_onboarding(
         and all onboarding tasks should allow run_task to specify additional
         or entirely override what's provided in a blueprint.
         """
+        shared_state = self.shared_state
+        assert isinstance(
+            shared_state, OnboardingSharedState
+        ), f"Cannot init onboarding config with {shared_state}, need OnboardingSharedState"
         data = onboarding_agent.state.get_data()
-        return self.shared_state.validate_onboarding(data)
+        return shared_state.validate_onboarding(data)