From 2ed483b5592de177ef93b0af23b394472c00d4e0 Mon Sep 17 00:00:00 2001
From: Tom Cobb <tom.cobb@diamond.ac.uk>
Date: Tue, 16 May 2023 13:12:50 +0000
Subject: [PATCH] Changes to work with a real PandA

- Block names can have numbers and underscores in
- Type hints seem to work better on the class, not instance
- Added demo
- Added type hints function which doesn't look through parent classes
---
 docs/examples/panda_demo.py      | 14 ++++++++++++++
 src/ophyd_epics_devices/panda.py | 20 +++++++++++++-------
 src/ophyd_epics_devices/utils.py | 13 +++++++++++++
 tests/test_utils.py              | 21 +++++++++++++++++++++
 4 files changed, 61 insertions(+), 7 deletions(-)
 create mode 100644 docs/examples/panda_demo.py
 create mode 100644 src/ophyd_epics_devices/utils.py
 create mode 100644 tests/test_utils.py

diff --git a/docs/examples/panda_demo.py b/docs/examples/panda_demo.py
new file mode 100644
index 0000000..0d32690
--- /dev/null
+++ b/docs/examples/panda_demo.py
@@ -0,0 +1,14 @@
+from bluesky import RunEngine
+
+# these three lines just let you use await statements
+# #in ipython terminal with the Run Engine event loop.
+from IPython import get_ipython
+from ophyd.v2.core import DeviceCollector
+
+from ophyd_epics_devices.panda import PandA
+
+get_ipython().run_line_magic("autoawait", "call_in_bluesky_event_loop")
+RE = RunEngine()
+
+with DeviceCollector():
+    my_panda = PandA("TS-PANDA")
diff --git a/src/ophyd_epics_devices/panda.py b/src/ophyd_epics_devices/panda.py
index 4f75bc0..9392bae 100644
--- a/src/ophyd_epics_devices/panda.py
+++ b/src/ophyd_epics_devices/panda.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import atexit
 import re
 from enum import Enum
@@ -34,6 +36,10 @@
 )
 from p4p.client.thread import Context
 
+from ophyd_epics_devices.utils import get_type_hints_no_inheritance
+
+ctxt = Context("pva")
+
 
 class PulseBlock(Device):
     delay: SignalRW[float]
@@ -93,7 +99,7 @@ class PVIEntry(TypedDict, total=False):
 
 
 def block_name_number(block_name: str) -> Tuple[str, int]:
-    m = re.match("^([a-z]+)([0-9]*)$", block_name)
+    m = re.match("^([0-9_a-z]+)([0-9]*)$", block_name)
     assert m, f"Expected '<block_name><block_num>', got '{block_name}'"
     name, num = m.groups()
     return name, int(num or 1)
@@ -144,7 +150,7 @@ def _del_ctxt():
 
     def verify_block(self, name: str, num: int):
         """Given a block name and number, return information about a block."""
-        anno = get_type_hints(self).get(name)
+        anno = get_type_hints(type(self)).get(name)
 
         block: Device = Device()
 
@@ -164,8 +170,7 @@ async def _make_block(self, name: str, num: int, block_pv: str, sim: bool = Fals
         sim mode then does a pvi call, and identifies this signal from the pvi call.
         """
         block = self.verify_block(name, num)
-
-        field_annos = get_type_hints(block)
+        field_annos = get_type_hints_no_inheritance(type(block))
         block_pvi = await pvi_get(block_pv, self.ctxt) if not sim else None
 
         # finds which fields this class actually has, e.g. delay, width...
@@ -231,7 +236,7 @@ def _make_signal(self, signal_pvi: PVIEntry, dtype: Optional[Type] = None):
         return signal_factory(dtype, "pva://" + read_pv, "pva://" + write_pv)
 
     def set_attribute(self, name, num, block):
-        anno = get_type_hints(self).get(name)
+        anno = get_type_hints(type(self)).get(name)
 
         # get_origin to see if it's a device vector.
         if (anno == DeviceVector[PulseBlock]) or (anno == DeviceVector[SeqBlock]):
@@ -251,10 +256,11 @@ async def connect(self, sim=False) -> None:
         pvi = await pvi_get(self._init_prefix + ":PVI", self.ctxt) if not sim else None
         hints = {
             attr_name: attr_type
-            for attr_name, attr_type in get_type_hints(self).items()
+            for attr_name, attr_type in get_type_hints_no_inheritance(
+                type(self)
+            ).items()
             if not attr_name.startswith("_")
         }
-
         # create all the blocks pvi says it should have,
         if pvi:
             for block_name, block_pvi in pvi.items():
diff --git a/src/ophyd_epics_devices/utils.py b/src/ophyd_epics_devices/utils.py
new file mode 100644
index 0000000..888f0a1
--- /dev/null
+++ b/src/ophyd_epics_devices/utils.py
@@ -0,0 +1,13 @@
+from typing import get_type_hints
+
+
+# Use with types, not instances
+def get_type_hints_no_inheritance(cls):
+    cls_hints = get_type_hints(cls)
+
+    for base_cls in cls.__bases__:
+        base_hints = get_type_hints(base_cls)
+        for base_hint_names in base_hints.keys():
+            if base_hint_names in cls_hints:
+                del cls_hints[base_hint_names]
+    return cls_hints
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..4d1353f
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,21 @@
+from typing import get_type_hints
+
+from ophyd_epics_devices.utils import get_type_hints_no_inheritance
+
+
+def test_get_type_hints_no_inheritance():
+    class BaseClass:
+        base_integer: int
+        base_string: str
+
+    class SubClass(BaseClass):
+        integer: int
+        string: str
+
+    assert get_type_hints(SubClass) == {
+        "base_integer": int,
+        "base_string": str,
+        "integer": int,
+        "string": str,
+    }
+    assert get_type_hints_no_inheritance(SubClass) == {"integer": int, "string": str}