From 3aefa644e515f658cadee5f48b755e98820e7103 Mon Sep 17 00:00:00 2001
From: Jan Petykiewicz <anewusername@gmail.com>
Date: Tue, 22 Jun 2021 18:25:31 -0700
Subject: [PATCH] fix type annotations for Device.__getitem__

---
 masque/builder/devices.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/masque/builder/devices.py b/masque/builder/devices.py
index a7d1e2d..384c5a4 100644
--- a/masque/builder/devices.py
+++ b/masque/builder/devices.py
@@ -1,4 +1,5 @@
 from typing import Dict, Iterable, List, Tuple, Union, TypeVar, Any, Iterator, Optional, Sequence
+from typing import overload, KeysView, ValuesView
 import copy
 import warnings
 import traceback
@@ -196,7 +197,15 @@ class Device(Copyable, Mirrorable):
 
         self._dead = False
 
-    def __getitem__(self, key: Union[str, Iterable[str]]) -> numpy.ndarray:
+    @overload
+    def __getitem__(self, key: str) -> Port:
+        pass
+
+    @overload
+    def __getitem__(self, key: Union[List[str], Tuple[str], KeysView[str], ValuesView[str]]) -> Dict[str, Port]:
+        pass
+
+    def __getitem__(self, key: Union[str, Iterable[str]]) -> Union[Port, Dict[str, Port]]:
         """
         For convenience, ports can be read out using square brackets:
         - `device['A'] == Port((0, 0), 0)`