From ebd12ce7166440924d2081d72ce7b6f1fce122e6 Mon Sep 17 00:00:00 2001
From: brunocruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Sat, 16 Dec 2023 16:12:34 -0800
Subject: [PATCH] Allow indexing of registers by register number (#22)

---
 harp/reader.py | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/harp/reader.py b/harp/reader.py
index deaa544..6b3645c 100644
--- a/harp/reader.py
+++ b/harp/reader.py
@@ -7,7 +7,8 @@
 from dataclasses import dataclass
 from numpy import dtype
 from pandas import DataFrame, Series
-from typing import Any, BinaryIO, Callable, Iterable, Optional, Protocol, Union
+from typing import Any, BinaryIO, Callable, Iterable, Mapping, Optional, Protocol, Union
+from collections import UserDict
 from pandas._typing import Axes
 from harp.model import BitMask, GroupMask, Model, PayloadMember, Register
 from harp.io import MessageType, read
@@ -44,13 +45,29 @@ def __init__(
         self.read = read
 
 
+class RegisterMap(UserDict[str, RegisterReader]):
+    _address_map: Mapping[int, RegisterReader]
+
+    def __init__(self, registers: Mapping[str, RegisterReader]) -> None:
+        super().__init__(registers)
+        self._address_map = {
+            value.register.address: value for value in registers.values()
+        }
+
+    def __getitem__(self, __key: Union[str, int]) -> RegisterReader:
+        if isinstance(__key, int):
+            return self._address_map[__key]
+        else:
+            return super().__getitem__(__key)
+
+
 class DeviceReader:
     device: Model
-    registers: dict[str, RegisterReader]
+    registers: RegisterMap
 
-    def __init__(self, device: Model, registers: dict[str, RegisterReader]) -> None:
+    def __init__(self, device: Model, registers: Mapping[str, RegisterReader]) -> None:
         self.device = device
-        self.registers = registers
+        self.registers = RegisterMap(registers)
 
     def __dir__(self) -> Iterable[str]:
         return self.registers.keys()