From 5d4020c127b938c2ba9648da965e97e63ac57a57 Mon Sep 17 00:00:00 2001
From: Callum Forrester <callum.forrester@diamond.ac.uk>
Date: Tue, 1 Aug 2023 14:20:31 +0100
Subject: [PATCH] Fix mypy errors

---
 src/blueapi/core/context.py                |  3 +--
 src/blueapi/plans/plans.py                 |  1 -
 src/blueapi/plugins/data_writing.py        | 22 +++-------------------
 src/blueapi/plugins/data_writing_server.py |  3 +--
 src/blueapi/worker/task.py                 |  4 +---
 tests/plugins/file_writing_detector.py     | 12 ++++--------
 tests/plugins/test_data_writing.py         |  3 ---
 7 files changed, 10 insertions(+), 38 deletions(-)

diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py
index f8f1b9497..206f1f6f6 100644
--- a/src/blueapi/core/context.py
+++ b/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,
diff --git a/src/blueapi/plans/plans.py b/src/blueapi/plans/plans.py
index 014b1cb56..4006541ce 100644
--- a/src/blueapi/plans/plans.py
+++ b/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
diff --git a/src/blueapi/plugins/data_writing.py b/src/blueapi/plugins/data_writing.py
index 00bed6f5c..104380a9d 100644
--- a/src/blueapi/plugins/data_writing.py
+++ b/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
diff --git a/src/blueapi/plugins/data_writing_server.py b/src/blueapi/plugins/data_writing_server.py
index fef673e71..3e5e40d98 100644
--- a/src/blueapi/plugins/data_writing_server.py
+++ b/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
diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py
index ced771828..d1b4de07d 100644
--- a/src/blueapi/worker/task.py
+++ b/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__)
diff --git a/tests/plugins/file_writing_detector.py b/tests/plugins/file_writing_detector.py
index 2fa0a498d..b95aa63c9 100644
--- a/tests/plugins/file_writing_detector.py
+++ b/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()
diff --git a/tests/plugins/test_data_writing.py b/tests/plugins/test_data_writing.py
index e3c3cc491..d827ecdd9 100644
--- a/tests/plugins/test_data_writing.py
+++ b/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 (