Skip to content

Commit

Permalink
Add custom type alias for floats or arrays of floats
Browse files Browse the repository at this point in the history
Closes #10
  • Loading branch information
ZedThree committed Oct 25, 2023
1 parent 360fb2c commit d5e3974
Show file tree
Hide file tree
Showing 4 changed files with 25 additions and 4 deletions.
5 changes: 3 additions & 2 deletions freeqdsk/_fileutils.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@

import fortranformat as ff
import numpy as np
from numpy.typing import ArrayLike

from ._typing import FloatArray, ArrayLike


def f2s(f: float) -> str:
Expand Down Expand Up @@ -145,7 +146,7 @@ def write_1d(values: Iterable[Any], out: ChunkOutput) -> None:
out.newline()


def write_2d(values: ArrayLike, out: ChunkOutput) -> None:
def write_2d(values: FloatArray, out: ChunkOutput) -> None:
r"""
Writes a 2D array to a ChunkOutput file handle.
Expand Down
20 changes: 20 additions & 0 deletions freeqdsk/_typing.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
"""Some type aliases, mostly for hinting about arrays of
numbers. `numpy.typing.ArrayLike` covers anything that can be
converted to an array of any type, and we mostly want to restrict to
arrays of floats.
"""

from typing import Union, Any, List
import numpy as np

try:
from numpy.typing import NDArray

float_ = np.floating[Any]
FloatArray = NDArray[np.floating[Any]]
except ImportError:
float_ = float # type: ignore
FloatArray = np.ndarray # type: ignore

ArrayLike = Union[int, float, float_, FloatArray, List[float_]]
2 changes: 1 addition & 1 deletion freeqdsk/aeqdsk.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
from typing import Dict, Optional, TextIO, Union

import numpy as np
from numpy.typing import ArrayLike

from ._typing import ArrayLike
from ._fileutils import read_array, write_array, write_line


Expand Down
2 changes: 1 addition & 1 deletion freeqdsk/geqdsk.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,8 +142,8 @@
from typing import Dict, Optional, TextIO, Union

import numpy as np
from numpy.typing import ArrayLike

from ._typing import ArrayLike
from ._fileutils import read_array, read_line, write_array, write_line


Expand Down

0 comments on commit d5e3974

Please sign in to comment.