From 2eaa5449bca55ac87475900dd64ca086bddc0023 Mon Sep 17 00:00:00 2001
From: Sebastiaan Huber <mail@sphuber.net>
Date: Tue, 29 Aug 2023 02:25:55 +0200
Subject: [PATCH] Add type hinting for `aiida.orm.nodes.data.array.xy`

---
 aiida/orm/nodes/data/array/xy.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/aiida/orm/nodes/data/array/xy.py b/aiida/orm/nodes/data/array/xy.py
index 7827fc050f..93177a6cea 100644
--- a/aiida/orm/nodes/data/array/xy.py
+++ b/aiida/orm/nodes/data/array/xy.py
@@ -12,16 +12,23 @@
 collections of y-arrays bound to a single x-array, and the methods to operate
 on them.
 """
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
 import numpy as np
 
 from aiida.common.exceptions import NotExistent
 
 from .array import ArrayData
 
+if TYPE_CHECKING:
+    from numpy import ndarray
+
 __all__ = ('XyData',)
 
 
-def check_convert_single_to_tuple(item):
+def check_convert_single_to_tuple(item: Any | Sequence[Any]) -> Sequence[Any]:
     """
     Checks if the item is a list or tuple, and converts it to a list if it is
     not already a list or tuple
@@ -44,7 +51,7 @@ class XyData(ArrayData):
     """
 
     @staticmethod
-    def _arrayandname_validator(array, name, units):
+    def _arrayandname_validator(array: 'ndarray', name: str, units: str) -> None:
         """
         Validates that the array is an numpy.ndarray and that the name is
         of type str. Raises TypeError or ValueError if this not the case.
@@ -61,7 +68,7 @@ def _arrayandname_validator(array, name, units):
         if not isinstance(units, str):
             raise TypeError('The units must always be a str.')
 
-    def set_x(self, x_array, x_name, x_units):
+    def set_x(self, x_array: 'ndarray', x_name: str, x_units: str) -> None:
         """
         Sets the array and the name for the x values.
 
@@ -74,7 +81,9 @@ def set_x(self, x_array, x_name, x_units):
         self.base.attributes.set('x_units', x_units)
         self.set_array('x_array', x_array)
 
-    def set_y(self, y_arrays, y_names, y_units):
+    def set_y(
+        self, y_arrays: 'ndarray' | Sequence['ndarray'], y_names: str | Sequence[str], y_units: str | Sequence[str]
+    ) -> None:
         """
         Set array(s) for the y part of the dataset. Also checks if the
         x_array has already been set, and that, the shape of the y_arrays
@@ -110,7 +119,7 @@ def set_y(self, y_arrays, y_names, y_units):
         self.base.attributes.set('y_names', y_names)
         self.base.attributes.set('y_units', y_units)
 
-    def get_x(self):
+    def get_x(self) -> tuple[str, 'ndarray', str]:
         """
         Tries to retrieve the x array and x name raises a NotExistent
         exception if no x array has been set yet.
@@ -126,7 +135,7 @@ def get_x(self):
             raise NotExistent('No x array has been set yet!')
         return x_name, x_array, x_units
 
-    def get_y(self):
+    def get_y(self) -> list[tuple[str, 'ndarray', str]]:
         """
         Tries to retrieve the y arrays and the y names, raises a
         NotExistent exception if they have not been set yet, or cannot be