From b9f65074983a1487d65bc0e2fb5a5736706fe6c2 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Sun, 24 Oct 2021 19:06:15 -0700 Subject: [PATCH] doc and type hint fixes --- gridlock/draw.py | 2 +- gridlock/grid.py | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/gridlock/draw.py b/gridlock/draw.py index 1142528..ca00e6b 100644 --- a/gridlock/draw.py +++ b/gridlock/draw.py @@ -198,7 +198,7 @@ def draw_polygon(self, surface_normal: Axis normal to the plane we're drawing on. Integer in `range(3)`. center: 3-element ndarray or list specifying an offset applied to the polygon polygon: Nx2 or Nx3 ndarray specifying the vertices of a polygon (non-closed, - clockwise). If Nx3, the surface_normal coordinate is ignored. Must have at + clockwise). If Nx3, the `surface_normal` coordinate is ignored. Must have at least 3 vertices. thickness: Thickness of the layer to draw eps: Value to draw with ('epsilon'). See `draw_polygons()` for details. diff --git a/gridlock/grid.py b/gridlock/grid.py index 9e828d1..9792858 100644 --- a/gridlock/grid.py +++ b/gridlock/grid.py @@ -1,4 +1,4 @@ -from typing import List, Tuple, Callable, Dict, Optional, Union, Sequence, ClassVar +from typing import List, Tuple, Callable, Dict, Optional, Union, Sequence, ClassVar, TypeVar import numpy # type: ignore from numpy import diff, floor, ceil, zeros, hstack, newaxis @@ -10,9 +10,8 @@ import copy from . import GridError -__author__ = 'Jan Petykiewicz' - eps_callable_type = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray] +T = TypeVar('T', bound='Grid') class Grid: @@ -227,7 +226,7 @@ class Grid: dxyz = self.shifted_dxyz(which_shifts) return [exyz[a][:-1] + dxyz[a] / 2.0 for a in range(3)] - def autoshifted_dxyz(self): + def autoshifted_dxyz(self) -> List[numpy.ndarray]: """ Return cell widths, with each dimension shifted by the corresponding shifts. @@ -325,7 +324,7 @@ class Grid: with open(filename, 'wb') as f: pickle.dump(self.__dict__, f, protocol=2) - def copy(self): + def copy(self: T) -> T: """ Returns: Deep copy of the grid.