Add type hints (#131)

addisonElliott · web-flow · commit f67ed36974e1 · 2022-09-19T20:18:00.000-05:00
diff --git a/README.rst b/README.rst
@@ -30,6 +30,11 @@ from numpy arrays.
 
 Dependencies
 ------------
+
+* `Numpy <https://numpy.org/>`_
+* nptyping
+* typing_extensions
+
 The module's only dependency is `numpy <http://numpy.scipy.org/>`_.
 
 Installation
diff --git a/nrrd/__init__.py b/nrrd/__init__.py
@@ -2,9 +2,10 @@
 from nrrd.formatters import *
 from nrrd.parsers import *
 from nrrd.reader import read, read_data, read_header
+from nrrd.types import NRRDFieldMap, NRRDFieldType, NRRDHeader
 from nrrd.writer import write
 
 __all__ = ['read', 'read_data', 'read_header', 'write', 'format_number_list', 'format_number', 'format_matrix',
            'format_optional_matrix', 'format_optional_vector', 'format_vector', 'parse_matrix',
            'parse_number_auto_dtype', 'parse_number_list', 'parse_optional_matrix', 'parse_optional_vector',
-           'parse_vector', '__version__']
+           'parse_vector', 'NRRDFieldType', 'NRRDFieldMap', 'NRRDHeader', '__version__']
diff --git a/nrrd/formatters.py b/nrrd/formatters.py
@@ -1,7 +1,11 @@
+from typing import Any, Optional, Union
+
+import nptyping as npt
 import numpy as np
+from typing_extensions import Literal
 
 
-def format_number(x):
+def format_number(x: Union[int, float]) -> str:
     """Format number to string
 
     Function converts a number to string. For numbers of class :class:`float`, up to 17 digits will be used to print
@@ -38,7 +42,7 @@ def format_number(x):
     return value
 
 
-def format_vector(x):
+def format_vector(x: npt.NDArray[Literal['*'], Any]) -> str:
     """Format a (N,) :class:`numpy.ndarray` into a NRRD vector string
 
     See :ref:`user-guide:int vector` and :ref:`user-guide:double vector` for more information on the format.
@@ -57,7 +61,7 @@ def format_vector(x):
     return '(' + ','.join([format_number(y) for y in x]) + ')'
 
 
-def format_optional_vector(x):
+def format_optional_vector(x: Optional[npt.NDArray[Literal['*'], Any]]) -> str:
     """Format a (N,) :class:`numpy.ndarray` into a NRRD optional vector string
 
     Function converts a (N,) :class:`numpy.ndarray` or :obj:`None` into a string using NRRD vector format. If the input
@@ -84,7 +88,7 @@ def format_optional_vector(x):
         return format_vector(x)
 
 
-def format_matrix(x):
+def format_matrix(x: npt.NDArray[Literal['*, *'], Any]) -> str:
     """Format a (M,N) :class:`numpy.ndarray` into a NRRD matrix string
 
     See :ref:`user-guide:int matrix` and :ref:`user-guide:double matrix` for more information on the format.
@@ -103,7 +107,7 @@ def format_matrix(x):
     return ' '.join([format_vector(y) for y in x])
 
 
-def format_optional_matrix(x):
+def format_optional_matrix(x: Optional[npt.NDArray[Literal['*, *'], Any]]) -> str:
     """Format a (M,N) :class:`numpy.ndarray` of :class:`float` into a NRRD optional matrix string
 
     Function converts a (M,N) :class:`numpy.ndarray` of :class:`float` into a string using the NRRD matrix format. For
@@ -129,7 +133,7 @@ def format_optional_matrix(x):
     return ' '.join([format_optional_vector(y) for y in x])
 
 
-def format_number_list(x):
+def format_number_list(x: npt.NDArray[Literal['*'], Any]) -> str:
     """Format a (N,) :class:`numpy.ndarray` into a NRRD number list.
 
     See :ref:`user-guide:int list` and :ref:`user-guide:double list` for more information on the format.
diff --git a/nrrd/parsers.py b/nrrd/parsers.py
@@ -1,9 +1,13 @@
+from typing import Any, Optional, Type, Union
+
+import nptyping as npt
 import numpy as np
+from typing_extensions import Literal
 
 from nrrd.errors import NRRDError
 
 
-def parse_vector(x, dtype=None):
+def parse_vector(x: str, dtype: Optional[Type[Union[int, float]]] = None) -> npt.NDArray[Literal['*'], Any]:
     """Parse NRRD vector from string into (N,) :class:`numpy.ndarray`.
 
     See :ref:`user-guide:int vector` and :ref:`user-guide:double vector` for more information on the format.
@@ -46,7 +50,8 @@ def parse_vector(x, dtype=None):
     return vector
 
 
-def parse_optional_vector(x, dtype=None):
+def parse_optional_vector(x: str, dtype: Optional[Type[Union[int, float]]] = None) -> \
+        Optional[npt.NDArray[Literal['*'], Any]]:
     """Parse optional NRRD vector from string into (N,) :class:`numpy.ndarray` or :obj:`None`.
 
     Function parses optional NRRD vector from string into an (N,) :class:`numpy.ndarray`. This function works the same
@@ -76,7 +81,7 @@ def parse_optional_vector(x, dtype=None):
         return parse_vector(x, dtype)
 
 
-def parse_matrix(x, dtype=None):
+def parse_matrix(x: str, dtype: Optional[Type[Union[int, float]]] = None) -> npt.NDArray[Literal['*, *'], Any]:
     """Parse NRRD matrix from string into (M,N) :class:`numpy.ndarray`.
 
     See :ref:`user-guide:int matrix` and :ref:`user-guide:double matrix` for more information on the format.
@@ -122,7 +127,7 @@ def parse_matrix(x, dtype=None):
     return matrix
 
 
-def parse_optional_matrix(x):
+def parse_optional_matrix(x: str) -> Optional[npt.NDArray[Literal['*, *'], Any]]:
     """Parse optional NRRD matrix from string into (M,N) :class:`numpy.ndarray` of :class:`float`.
 
     Function parses optional NRRD matrix from string into an (M,N) :class:`numpy.ndarray` of :class:`float`. This
@@ -165,7 +170,7 @@ def parse_optional_matrix(x):
     return matrix
 
 
-def parse_number_list(x, dtype=None):
+def parse_number_list(x: str, dtype: Optional[Type[Union[int, float]]] = None) -> npt.NDArray[Literal['*'], Any]:
     """Parse NRRD number list from string into (N,) :class:`numpy.ndarray`.
 
     See :ref:`user-guide:int list` and :ref:`user-guide:double list` for more information on the format.
@@ -204,7 +209,7 @@ def parse_number_list(x, dtype=None):
     return number_list
 
 
-def parse_number_auto_dtype(x):
+def parse_number_auto_dtype(x: str) -> Union[int, float]:
     """Parse number from string with automatic type detection.
 
     Parses input string and converts to a number using automatic type detection. If the number contains any
@@ -223,7 +228,7 @@ def parse_number_auto_dtype(x):
         Number parsed from :obj:`x` string
     """
 
-    value = float(x)
+    value: Union[int, float] = float(x)
 
     if value.is_integer():
         value = int(value)
diff --git a/nrrd/reader.py b/nrrd/reader.py
@@ -6,14 +6,16 @@
 import warnings
 import zlib
 from collections import OrderedDict
+from typing import IO, AnyStr, Iterable, Tuple
 
 from nrrd.parsers import *
+from nrrd.types import IndexOrder, NRRDFieldMap, NRRDFieldType, NRRDHeader
 
 # Older versions of Python had issues when uncompressed data was larger than 4GB (2^32). This should be fixed in latest
 # version of Python 2.7 and all versions of Python 3. The fix for this issue is to read the data in smaller chunks.
 # Chunk size is set to be large at 1GB to improve performance. If issues arise decompressing larger files, try to reduce
 # this value
-_READ_CHUNKSIZE = 2 ** 32
+_READ_CHUNKSIZE: int = 2 ** 32
 
 _NRRD_REQUIRED_FIELDS = ['dimension', 'type', 'encoding', 'sizes']
 
@@ -84,7 +86,7 @@
 }
 
 
-def _get_field_type(field, custom_field_map):
+def _get_field_type(field: str, custom_field_map: Optional[NRRDFieldMap]) -> NRRDFieldType:
     if field in ['dimension', 'lineskip', 'line skip', 'byteskip', 'byte skip', 'space dimension']:
         return 'int'
     elif field in ['min', 'max', 'oldmin', 'old min', 'oldmax', 'old max']:
@@ -99,7 +101,7 @@ def _get_field_type(field, custom_field_map):
         return 'string list'
     elif field in ['labels', 'units', 'space units']:
         return 'quoted string list'
-    # No int vector fields as of now
+    # No int vector fields yet
     # elif field in []:
     #     return 'int vector'
     elif field in ['space origin']:
@@ -116,7 +118,7 @@ def _get_field_type(field, custom_field_map):
         return 'string'
 
 
-def _parse_field_value(value, field_type):
+def _parse_field_value(value: str, field_type: NRRDFieldType) -> Any:
     if field_type == 'int':
         return int(value)
     elif field_type == 'double':
@@ -146,28 +148,28 @@ def _parse_field_value(value, field_type):
         raise NRRDError(f'Invalid field type given: {field_type}')
 
 
-def _determine_datatype(fields):
+def _determine_datatype(header: NRRDHeader) -> np.dtype:
     """Determine the numpy dtype of the data."""
 
     # Convert the NRRD type string identifier into a NumPy string identifier using a map
-    np_typestring = _TYPEMAP_NRRD2NUMPY[fields['type']]
+    np_typestring = _TYPEMAP_NRRD2NUMPY[header['type']]
 
     # This is only added if the datatype has more than one byte and is not using ASCII encoding
     # Note: Endian is not required for ASCII encoding
-    if np.dtype(np_typestring).itemsize > 1 and fields['encoding'] not in ['ASCII', 'ascii', 'text', 'txt']:
-        if 'endian' not in fields:
+    if np.dtype(np_typestring).itemsize > 1 and header['encoding'] not in ['ASCII', 'ascii', 'text', 'txt']:
+        if 'endian' not in header:
             raise NRRDError('Header is missing required field: endian')
-        elif fields['endian'] == 'big':
+        elif header['endian'] == 'big':
             np_typestring = '>' + np_typestring
-        elif fields['endian'] == 'little':
+        elif header['endian'] == 'little':
             np_typestring = '<' + np_typestring
         else:
-            raise NRRDError(f'Invalid endian value in header: {fields["endian"]}')
+            raise NRRDError(f'Invalid endian value in header: {header["endian"]}')
 
     return np.dtype(np_typestring)
 
 
-def _validate_magic_line(line):
+def _validate_magic_line(line: str) -> int:
     """For NRRD files, the first four characters are always "NRRD", and
     remaining characters give information about the file format version
 
@@ -197,7 +199,7 @@ def _validate_magic_line(line):
     return len(line)
 
 
-def read_header(file, custom_field_map=None):
+def read_header(file: Union[str, Iterable[AnyStr]], custom_field_map: Optional[NRRDFieldMap] = None) -> NRRDHeader:
     """Read contents of header and parse values from :obj:`file`
 
     :obj:`file` can be a filename indicating where the NRRD header is located or a string iterator object. If a
@@ -284,7 +286,7 @@ def read_header(file, custom_field_map=None):
             else:
                 warnings.warn(f'Duplicate header field: {field}')
 
-        # Get the datatype of the field based on it's field name and custom field map
+        # Get the datatype of the field based on its field name and custom field map
         field_type = _get_field_type(field, custom_field_map)
 
         # Parse the field value using the datatype retrieved
@@ -299,7 +301,8 @@ def read_header(file, custom_field_map=None):
     return header
 
 
-def read_data(header, fh=None, filename=None, index_order='F'):
+def read_data(header: NRRDHeader, fh: Optional[IO] = None, filename: Optional[str] = None,
+              index_order: IndexOrder = 'F') -> npt.NDArray:
     """Read data from file into :class:`numpy.ndarray`
 
     The two parameters :obj:`fh` and :obj:`filename` are optional depending on the parameters but it never hurts to
@@ -427,7 +430,7 @@ def read_data(header, fh=None, filename=None, index_order='F'):
         # Loop through the file and read a chunk at a time (see _READ_CHUNKSIZE why it is read in chunks)
         decompressed_data = bytearray()
 
-        # Read all of the remaining data from the file
+        # Read all the remaining data from the file
         # Obtain the length of the compressed data since we will be using it repeatedly, more efficient
         compressed_data = fh.read()
         compressed_data_len = len(compressed_data)
@@ -474,7 +477,8 @@ def read_data(header, fh=None, filename=None, index_order='F'):
     return data
 
 
-def read(filename, custom_field_map=None, index_order='F'):
+def read(filename: str, custom_field_map: Optional[NRRDFieldMap] = None, index_order: IndexOrder = 'F') \
+        -> Tuple[npt.NDArray, NRRDHeader]:
     """Read a NRRD file and return the header and data
 
     See :ref:`user-guide:Reading NRRD files` for more information on reading NRRD files.
diff --git a/nrrd/types.py b/nrrd/types.py
@@ -0,0 +1,11 @@
+from typing import Any, Dict
+
+from typing_extensions import Literal
+
+NRRDFieldType = Literal['int', 'double', 'string', 'int list', 'double list', 'string list', 'quoted string list',
+                        'int vector', 'double vector', 'int matrix', 'double matrix']
+
+IndexOrder = Literal['F', 'C']
+
+NRRDFieldMap = Dict[str, NRRDFieldType]
+NRRDHeader = Dict[str, Any]
diff --git a/nrrd/writer.py b/nrrd/writer.py
@@ -4,16 +4,20 @@
 import zlib
 from collections import OrderedDict
 from datetime import datetime
+from typing import IO, Dict
+
+import nptyping as npt
 
 from nrrd.errors import NRRDError
 from nrrd.formatters import *
 from nrrd.reader import _get_field_type
+from nrrd.types import IndexOrder, NRRDFieldMap, NRRDFieldType, NRRDHeader
 
 # Older versions of Python had issues when uncompressed data was larger than 4GB (2^32). This should be fixed in latest
 # version of Python 2.7 and all versions of Python 3. The fix for this issue is to read the data in smaller chunks. The
 # chunk size is set to be small here at 1MB since performance did not vary much based on the chunk size. A smaller chunk
 # size has the benefit of using less RAM at once.
-_WRITE_CHUNKSIZE = 2 ** 20
+_WRITE_CHUNKSIZE: int = 2 ** 20
 
 _NRRD_FIELD_ORDER = [
     'type',
@@ -45,7 +49,8 @@
     'space units',
     'space origin',
     'measurement frame',
-    'data file']
+    'data file'
+]
 
 _TYPEMAP_NUMPY2NRRD = {
     'i1': 'int8',
@@ -69,7 +74,7 @@
 }
 
 
-def _format_field_value(value, field_type):
+def _format_field_value(value: Any, field_type: NRRDFieldType) -> str:
     if field_type == 'int':
         return format_number(value)
     elif field_type == 'double':
@@ -96,7 +101,7 @@ def _format_field_value(value, field_type):
         raise NRRDError(f'Invalid field type given: {field_type}')
 
 
-def _handle_header(data, header=None, index_order='F'):
+def _handle_header(data: npt.NDArray, header: Optional[NRRDHeader] = None, index_order: IndexOrder = 'F') -> NRRDHeader:
     if header is None:
         header = {}
 
@@ -133,7 +138,7 @@ def _handle_header(data, header=None, index_order='F'):
     return header
 
 
-def _write_header(file, header, custom_field_map=None):
+def _write_header(file: IO, header: Dict[str, Any], custom_field_map: Optional[NRRDFieldMap] = None):
     file.write(b'NRRD0005\n')
     file.write(b'# This NRRD file was generated by pynrrd\n')
     file.write(b'# on ' + datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S').encode('ascii') + b'(GMT).\n')
@@ -177,7 +182,8 @@ def _write_header(file, header, custom_field_map=None):
     file.write(b'\n')
 
 
-def _write_data(data, fh, header, compression_level=None, index_order='F'):
+def _write_data(data: npt.NDArray, fh: IO, header: NRRDHeader, compression_level: Optional[int] = None,
+                index_order: IndexOrder = 'F'):
     if index_order not in ['F', 'C']:
         raise NRRDError('Invalid index order')
 
@@ -243,8 +249,9 @@ def _write_data(data, fh, header, compression_level=None, index_order='F'):
         fh.flush()
 
 
-def write(file, data, header=None, detached_header=False, relative_data_path=True, custom_field_map=None,
-          compression_level=9, index_order='F'):
+def write(file: Union[str, IO], data: npt.NDArray, header: Optional[NRRDHeader] = None,
+          detached_header: bool = False, relative_data_path: bool = True,
+          custom_field_map: Optional[NRRDFieldMap] = None, compression_level: int = 9, index_order: IndexOrder = 'F'):
     """Write :class:`numpy.ndarray` to NRRD file
 
     The :obj:`file` parameter specifies the absolute or relative filename to write the NRRD file to or an
@@ -346,14 +353,12 @@ def write(file, data, header=None, detached_header=False, relative_data_path=Tru
         # Update the data file field in the header with the path of the detached data
         # TODO This will cause problems when the user specifies a relative data path and gives a custom path OUTSIDE
         #  of the current directory.
-        header['data file'] = os.path.basename(data_filename) \
-            if relative_data_path else os.path.abspath(data_filename)
+        header['data file'] = os.path.basename(data_filename) if relative_data_path else os.path.abspath(data_filename)
         detached_header = True
     elif file.endswith('.nrrd') and detached_header:
         data_filename = file
         file = f'{os.path.splitext(file)[0]}.nhdr'
-        header['data file'] = os.path.basename(data_filename) \
-            if relative_data_path else os.path.abspath(data_filename)
+        header['data file'] = os.path.basename(data_filename) if relative_data_path else os.path.abspath(data_filename)
     else:
         # Write header & data as one file
         data_filename = file
diff --git a/requirements.txt b/requirements.txt
@@ -1,2 +1,3 @@
 numpy>=1.11.1
+nptyping
 typing_extensions

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
`1`	`1`	`numpy>=1.11.1`
	`2`	`+nptyping`
`2`	`3`	`typing_extensions`