From d5fbcbb75979ab2109582f5fb67f4b1030f96c68 Mon Sep 17 00:00:00 2001
From: Tobias Ahrens <tobias.ahrens@zhinst.com>
Date: Thu, 10 Aug 2023 21:28:54 +0200
Subject: [PATCH] Use numpy to check for interger type in NodeList (#252)

The getitem function of NodeList differentiates between integers and nodes as
argument. This guarantees a concistent behaviour. Wether the agrument is an
integer or not needs to be checked. This check now also supports all numpy
types and therefore is more versatile.
---
 CHANGELOG.md                        |  3 +++
 src/zhinst/toolkit/nodetree/node.py |  4 +++-
 tests/test_nodetree.py              | 11 ++++++++---
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 42cd4c0c..2bfc2bf1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,8 @@
 # zhinst-toolkit Changelog
 
+## Version 0.6.1
+* Fix #252. ListNodes now also supports numpy interger types as index argument.
+
 ## Version 0.6.0
 * Revert full support of `fnmatch` wildcards and instead use the LabOne wildcard support.
   This means only `*` symbols are supported. A `*` in the middle of the path matches
diff --git a/src/zhinst/toolkit/nodetree/node.py b/src/zhinst/toolkit/nodetree/node.py
index 91932445..27fe6545 100644
--- a/src/zhinst/toolkit/nodetree/node.py
+++ b/src/zhinst/toolkit/nodetree/node.py
@@ -9,6 +9,7 @@
 from collections import namedtuple
 from collections.abc import Sequence
 from enum import IntEnum
+import numpy as np
 
 from zhinst.toolkit.nodetree.helper import (
     NodeDict,
@@ -1127,7 +1128,8 @@ def __getitem__(
         ...
 
     def __getitem__(self, item):
-        if isinstance(item, int):
+        # User numpy check here to ensure numpy types are handled correctly (#252)
+        if np.issubdtype(type(item), np.integer):
             return self._elements[item]
         return Node(self._root, self._tree + (str(item),))
 
diff --git a/tests/test_nodetree.py b/tests/test_nodetree.py
index 9535b33d..bd9f28e4 100644
--- a/tests/test_nodetree.py
+++ b/tests/test_nodetree.py
@@ -9,7 +9,7 @@
 import gc
 
 import pytest
-from numpy import array as nparray
+import numpy as np
 
 from zhinst.toolkit.driver.devices import HDAWG
 from zhinst.toolkit.nodetree import Node, NodeTree
@@ -381,7 +381,7 @@ def test_get_deep(connection):
     assert data == array("l", [123, 10, 32, 10, 125, 10])
     # HF2 node
     connection.get.return_value = OrderedDict(
-        [("/dev1234/demods/0/rate", nparray([1674.10717773]))]
+        [("/dev1234/demods/0/rate", np.array([1674.10717773]))]
     )
     timestamp, data = tree.demods[0].sample(deep=True)
     assert data == 1674.10717773
@@ -424,7 +424,7 @@ def test_get_wildcard(connection):
 
     # HF2 support
     connection.get.return_value = OrderedDict(
-        [("/dev1234/demods/0/impedance", nparray([125]))]
+        [("/dev1234/demods/0/impedance", np.array([125]))]
     )
     result = tree.demods()
     assert result[tree.demods[0].impedance] == 125
@@ -1026,6 +1026,11 @@ def test_nodelist_is_node(connection, hdawg):
     assert bar[0] == hdawg
     assert bar == Node(nt, ("foobar",))
 
+def test_nodelist_numpy(connection, hdawg):
+    nt = NodeTree(connection, "DEV1234")
+    bar = NodeList([hdawg, hdawg], nt, ("foobar",))
+    for test_index in np.arange(2):
+        assert bar[test_index] == hdawg
 
 def test_nodelist_hash(connection, hdawg):
     nt = NodeTree(connection, "DEV1234")