Fix mypy errors

src/blueapi/core/context.py

@@ -9,7 +9,6 @@
     Callable,
     Dict,
     Generic,
-    Iterable,
     List,
     Optional,
     Sequence,
@@ -63,7 +62,7 @@ class BlueskyContext:
 
     _reference_cache: Dict[Type, Type] = field(default_factory=dict)
 
-    def wrap(self, plan: MsgGenerator) -> Iterable[PlanWrapper]:
+    def wrap(self, plan: MsgGenerator) -> MsgGenerator:
         wrapped_plan = functools.reduce(
             lambda wrapped, next_wrapper: next_wrapper(wrapped),
             self.plan_wrappers,

src/blueapi/plans/plans.py

@@ -5,7 +5,6 @@
 import bluesky.plans as bp
 from bluesky.protocols import Movable, Readable
 from cycler import Cycler, cycler
-from dodal.adsim import AdSimDetector
 from scanspec.specs import Spec
 
 from blueapi.core import MsgGenerator

src/blueapi/plugins/data_writing.py

@@ -1,29 +1,13 @@
 import itertools
 from abc import ABC, abstractmethod, abstractproperty
-from collections import deque
 from pathlib import Path
-from typing import (
-    Any,
-    Callable,
-    Deque,
-    Dict,
-    Generic,
-    Iterable,
-    List,
-    Optional,
-    Protocol,
-    TypeVar,
-    runtime_checkable,
-)
+from typing import Optional
 
 import bluesky.plan_stubs as bps
-import requests
 from aiohttp import ClientSession
-from bluesky.protocols import Movable
-from bluesky.utils import Msg, make_decorator
-from ophyd.areadetector.filestore_mixins import FileStoreBase
+from bluesky.utils import make_decorator
 
-from blueapi.core import BlueskyContext, Device, MsgGenerator, walk_devices
+from blueapi.core import MsgGenerator
 from blueapi.plugins.data_writing_server import DataCollection
 
 from .data_writing_server import DataCollection, DataCollectionSetupResult

src/blueapi/plugins/data_writing_server.py

@@ -1,8 +1,7 @@
 import itertools
 import os
-from enum import Enum
 from pathlib import Path
-from typing import Dict, Mapping
+from typing import Dict
 
 from fastapi import FastAPI
 from pydantic import BaseModel

src/blueapi/worker/task.py

@@ -1,11 +1,9 @@
-import functools
 import logging
-from typing import Any, Iterable, Mapping, Optional
+from typing import Any, Mapping, Optional
 
 from pydantic import BaseModel, Field
 
 from blueapi.core import BlueskyContext
-from blueapi.core.bluesky_types import MsgGenerator, PlanWrapper
 from blueapi.utils import BlueapiBaseModel
 
 LOGGER = logging.getLogger(__name__)

tests/plugins/file_writing_detector.py

@@ -1,12 +1,6 @@
-import itertools
-import uuid
-from pathlib import Path
-from typing import List, Optional, Tuple
+from typing import List
 
-import h5py as h5
-import numpy as np
-from ophyd import Component, Device, Signal, SignalRO
-from ophyd.sim import EnumSignal, SynGauss, SynSignal, SynSignalRO
+from ophyd import Component, Device, Signal
 
 from blueapi.plugins.data_writing import DataCollectionProvider
 
@@ -27,5 +21,7 @@ def trigger(self, *args, **kwargs):
 
     def stage(self) -> List[object]:
         collection = self._provider.current_data_collection
+        if collection is None:
+            raise Exception("No active collection")
         self.stage_sigs[self.collection_number] = collection.collection_number
         return super().stage()

tests/plugins/test_data_writing.py

@@ -9,11 +9,8 @@
     run_wrapper,
     set_run_key_decorator,
     set_run_key_wrapper,
-    stage_decorator,
     stage_wrapper,
 )
-from bluesky.protocols import Readable
-from ophyd.sim import SynAxis
 
 from blueapi.core import DataEvent, MsgGenerator
 from blueapi.plugins.data_writing import (