fix type errors

Author: aaxelb

addon_service/addon_operation/models.py

@@ -7,7 +7,10 @@
     get_imp_name,
 )
 from addon_service.common.static_dataclass_model import StaticDataclassModel
-from addon_toolkit import AddonOperationImp
+from addon_toolkit import (
+    AddonCapabilities,
+    AddonOperationImp,
+)
 from addon_toolkit.json_arguments import jsonschema_for_signature_params
 from addon_toolkit.operation import AddonOperationType
 
@@ -42,7 +45,7 @@ def implementation_docstring(self) -> str:
         return self.operation_imp.imp_function.__doc__ or ""
 
     @cached_property
-    def capability(self) -> str:
+    def capability(self) -> AddonCapabilities:
         return self.operation_imp.declaration.capability
 
     @cached_property

addon_service/addon_operation/serializers.py

@@ -1,11 +1,11 @@
 from rest_framework_json_api import serializers
 from rest_framework_json_api.utils import get_resource_type_from_model
 
-from addon_service.addon_imp.models import AddonImp
 from addon_service.common import view_names
 from addon_service.common.enums.serializers import EnumNameChoiceField
 from addon_service.common.serializer_fields import DataclassRelatedDataField
 from addon_toolkit import AddonCapabilities
+from addon_toolkit.imp import AddonImp
 
 from .models import AddonOperationModel
 

addon_toolkit/operation.py

@@ -7,6 +7,7 @@
     Iterator,
 )
 
+from .capabilities import AddonCapabilities
 from .declarator import Declarator
 
 
@@ -35,7 +36,7 @@ class AddonOperationDeclaration:
     """
 
     operation_type: AddonOperationType
-    capability: enum.Enum
+    capability: AddonCapabilities
     operation_fn: Callable  # the decorated function
     return_type: type = dataclasses.field(
         default=type(None),  # if not provided, inferred by __post_init__

addon_toolkit/tests/test_addon_protocol.py

@@ -21,7 +21,7 @@ class TestAddonProtocol(unittest.TestCase):
     # the basics of an addon protocol
 
     ###
-    # shared test env (on `self`)
+    # shared test env (initialized by setUpClass)
     _MyProtocol: type  # typing.Protocol subclass decorated with `@addon_protocol`
     _MyImplementation: type  # subclass of _MyProtocol
     _my_imp: AddonImp