From 2bc03cbbf4f09c9fc4deaf4f2323c9a62437d98a Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Sun, 1 Nov 2020 19:33:43 -0800 Subject: [PATCH] improve some more type annotations using TypeVar --- masque/file/klamath.py | 1 + masque/label.py | 17 ++++++---- masque/pattern.py | 72 ++++++++++++++++++++++-------------------- masque/subpattern.py | 19 ++++++----- 4 files changed, 60 insertions(+), 49 deletions(-) diff --git a/masque/file/klamath.py b/masque/file/klamath.py index 0f858cf..1c22e73 100644 --- a/masque/file/klamath.py +++ b/masque/file/klamath.py @@ -576,6 +576,7 @@ def load_library(stream: BinaryIO, if is_secondary is None: def is_secondary(k: str): return False + assert(is_secondary is not None) stream.seek(0) library_info = _read_header(stream) diff --git a/masque/label.py b/masque/label.py index 5027af5..b436653 100644 --- a/masque/label.py +++ b/masque/label.py @@ -1,4 +1,4 @@ -from typing import Tuple, Dict, Optional +from typing import Tuple, Dict, Optional, TypeVar import copy import numpy # type: ignore @@ -8,6 +8,9 @@ from .traits import PositionableImpl, LayerableImpl, Copyable, Pivotable, Lockab from .traits import AnnotatableImpl +L = TypeVar('L', bound='Label') + + class Label(PositionableImpl, LayerableImpl, LockableImpl, RepeatableImpl, AnnotatableImpl, Pivotable, Copyable, metaclass=AutoSlots): """ @@ -44,7 +47,7 @@ class Label(PositionableImpl, LayerableImpl, LockableImpl, RepeatableImpl, Annot repetition: Optional[Repetition] = None, annotations: Optional[annotations_t] = None, locked: bool = False, - ): + ) -> None: LockableImpl.unlock(self) self.identifier = () self.string = string @@ -54,21 +57,21 @@ class Label(PositionableImpl, LayerableImpl, LockableImpl, RepeatableImpl, Annot self.annotations = annotations if annotations is not None else {} self.set_locked(locked) - def __copy__(self) -> 'Label': + def __copy__(self: L) -> L: return Label(string=self.string, offset=self.offset.copy(), layer=self.layer, repetition=self.repetition, locked=self.locked) - def __deepcopy__(self, memo: Dict = None) -> 'Label': + def __deepcopy__(self: L, memo: Dict = None) -> L: memo = {} if memo is None else memo new = copy.copy(self).unlock() new._offset = self._offset.copy() new.set_locked(self.locked) return new - def rotate_around(self, pivot: vector2, rotation: float) -> 'Label': + def rotate_around(self: L, pivot: vector2, rotation: float) -> L: """ Rotate the label around a point. @@ -98,12 +101,12 @@ class Label(PositionableImpl, LayerableImpl, LockableImpl, RepeatableImpl, Annot """ return numpy.array([self.offset, self.offset]) - def lock(self) -> 'Label': + def lock(self: L) -> L: PositionableImpl._lock(self) LockableImpl.lock(self) return self - def unlock(self) -> 'Label': + def unlock(self: L) -> L: LockableImpl.unlock(self) PositionableImpl._unlock(self) return self diff --git a/masque/pattern.py b/masque/pattern.py index 33c5030..fd1b8de 100644 --- a/masque/pattern.py +++ b/masque/pattern.py @@ -3,7 +3,7 @@ """ from typing import List, Callable, Tuple, Dict, Union, Set, Sequence, Optional, Type, overload -from typing import MutableMapping, Iterable +from typing import MutableMapping, Iterable, TypeVar, Any import copy import pickle from itertools import chain @@ -24,6 +24,9 @@ from .traits import LockableImpl, AnnotatableImpl, Scalable visitor_function_t = Callable[['Pattern', Tuple['Pattern'], Dict, numpy.ndarray], 'Pattern'] +P = TypeVar('P', bound='Pattern') + + class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): """ 2D layout consisting of some set of shapes, labels, and references to other Pattern objects @@ -56,7 +59,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): subpatterns: Sequence[SubPattern] = (), annotations: Optional[annotations_t] = None, locked: bool = False, - ): + ) -> None: """ Basic init; arguments get assigned to member variables. Non-list inputs for shapes and subpatterns get converted to lists. @@ -112,11 +115,11 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): locked=self.locked) return new - def rename(self, name: str) -> 'Pattern': + def rename(self: P, name: str) -> P: self.name = name return self - def append(self, other_pattern: 'Pattern') -> 'Pattern': + def append(self: P, other_pattern: P) -> P: """ Appends all shapes, labels and subpatterns from other_pattern to self's shapes, labels, and supbatterns. @@ -220,13 +223,13 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): pat = memo[pat_id] return pat - def dfs(self, + def dfs(self: P, visit_before: visitor_function_t = None, visit_after: visitor_function_t = None, transform: Union[numpy.ndarray, bool, None] = False, memo: Optional[Dict] = None, - hierarchy: Tuple['Pattern', ...] = (), - ) -> 'Pattern': + hierarchy: Tuple[P, ...] = (), + ) -> P: """ Experimental convenience function. Performs a depth-first traversal of this pattern and its subpatterns. @@ -305,10 +308,10 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): pat = visit_after(pat, hierarchy=hierarchy, memo=memo, transform=transform) # type: ignore return pat - def polygonize(self, + def polygonize(self: P, poly_num_points: Optional[int] = None, poly_max_arclen: Optional[float] = None, - ) -> 'Pattern': + ) -> P: """ Calls `.to_polygons(...)` on all the shapes in this Pattern and any referenced patterns, replacing them with the returned polygons. @@ -333,10 +336,10 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): subpat.pattern.polygonize(poly_num_points, poly_max_arclen) return self - def manhattanize(self, + def manhattanize(self: P, grid_x: numpy.ndarray, grid_y: numpy.ndarray, - ) -> 'Pattern': + ) -> P: """ Calls `.polygonize()` and `.flatten()` on the pattern, then calls `.manhattanize()` on all the resulting shapes, replacing them with the returned Manhattan polygons. @@ -355,11 +358,11 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): (shape.manhattanize(grid_x, grid_y) for shape in old_shapes))) return self - def subpatternize(self, + def subpatternize(self: P, recursive: bool = True, norm_value: int = int(1e6), exclude_types: Tuple[Type] = (Polygon,) - ) -> 'Pattern': + ) -> P: """ Iterates through this `Pattern` and all referenced `Pattern`s. Within each `Pattern`, it iterates over all shapes, calling `.normalized_form(norm_value)` on them to retrieve a scale-, @@ -476,7 +479,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): ids.update(pat.referenced_patterns_by_id()) return ids - def referenced_patterns_by_name(self, **kwargs) -> List[Tuple[Optional[str], Optional['Pattern']]]: + def referenced_patterns_by_name(self, **kwargs: Any) -> List[Tuple[Optional[str], Optional['Pattern']]]: """ Create a list of `(pat.name, pat)` tuples for all Pattern objects referenced by this Pattern (operates recursively on all referenced Patterns as well). @@ -544,7 +547,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): else: return numpy.vstack((min_bounds, max_bounds)) - def flatten(self) -> 'Pattern': + def flatten(self: P) -> P: """ Removes all subpatterns and adds equivalent shapes. @@ -580,10 +583,10 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): self.append(p) return self - def wrap_repeated_shapes(self, + def wrap_repeated_shapes(self: P, name_func: Callable[['Pattern', Union[Shape, Label]], str] = lambda p, s: '_repetition', recursive: bool = True, - ) -> 'Pattern': + ) -> P: """ Wraps all shapes and labels with a non-`None` `repetition` attribute into a `SubPattern`/`Pattern` combination, and applies the `repetition` @@ -625,7 +628,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): return self - def translate_elements(self, offset: vector2) -> 'Pattern': + def translate_elements(self: P, offset: vector2) -> P: """ Translates all shapes, label, and subpatterns by the given offset. @@ -639,7 +642,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): entry.translate(offset) return self - def scale_elements(self, c: float) -> 'Pattern': + def scale_elements(self: P, c: float) -> P: """" Scales all shapes and subpatterns by the given value. @@ -653,7 +656,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): entry.scale_by(c) return self - def scale_by(self, c: float) -> 'Pattern': + def scale_by(self: P, c: float) -> P: """ Scale this Pattern by the given value (all shapes and subpatterns and their offsets are scaled) @@ -672,7 +675,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): label.offset *= c return self - def rotate_around(self, pivot: vector2, rotation: float) -> 'Pattern': + def rotate_around(self: P, pivot: vector2, rotation: float) -> P: """ Rotate the Pattern around the a location. @@ -690,7 +693,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): self.translate_elements(+pivot) return self - def rotate_element_centers(self, rotation: float) -> 'Pattern': + def rotate_element_centers(self: P, rotation: float) -> P: """ Rotate the offsets of all shapes, labels, and subpatterns around (0, 0) @@ -704,7 +707,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): entry.offset = numpy.dot(rotation_matrix_2d(rotation), entry.offset) return self - def rotate_elements(self, rotation: float) -> 'Pattern': + def rotate_elements(self: P, rotation: float) -> P: """ Rotate each shape and subpattern around its center (offset) @@ -718,7 +721,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): entry.rotate(rotation) return self - def mirror_element_centers(self, axis: int) -> 'Pattern': + def mirror_element_centers(self: P, axis: int) -> P: """ Mirror the offsets of all shapes, labels, and subpatterns across an axis @@ -733,7 +736,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): entry.offset[axis - 1] *= -1 return self - def mirror_elements(self, axis: int) -> 'Pattern': + def mirror_elements(self: P, axis: int) -> P: """ Mirror each shape and subpattern across an axis, relative to its offset @@ -749,7 +752,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): entry.mirror(axis) return self - def mirror(self, axis: int) -> 'Pattern': + def mirror(self: P, axis: int) -> P: """ Mirror the Pattern across an axis @@ -764,7 +767,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): self.mirror_element_centers(axis) return self - def scale_element_doses(self, c: float) -> 'Pattern': + def scale_element_doses(self: P, c: float) -> P: """ Multiply all shape and subpattern doses by a factor @@ -778,7 +781,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): entry.dose *= c return self - def copy(self) -> 'Pattern': + def copy(self: P) -> P: """ Return a copy of the Pattern, deep-copying shapes and copying subpattern entries, but not deep-copying any referenced patterns. @@ -790,7 +793,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): """ return copy.copy(self) - def deepcopy(self) -> 'Pattern': + def deepcopy(self: P) -> P: """ Convenience method for `copy.deepcopy(pattern)` @@ -808,7 +811,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): and len(self.shapes) == 0 and len(self.labels) == 0) - def lock(self) -> 'Pattern': + def lock(self: P) -> P: """ Lock the pattern, raising an exception if it is modified. Also see `deeplock()`. @@ -823,7 +826,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): LockableImpl.lock(self) return self - def unlock(self) -> 'Pattern': + def unlock(self: P) -> P: """ Unlock the pattern @@ -837,7 +840,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): self.subpatterns = list(self.subpatterns) return self - def deeplock(self) -> 'Pattern': + def deeplock(self: P) -> P: """ Recursively lock the pattern, all referenced shapes, subpatterns, and labels. @@ -851,7 +854,7 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): sp.deeplock() return self - def deepunlock(self) -> 'Pattern': + def deepunlock(self: P) -> P: """ Recursively unlock the pattern, all referenced shapes, subpatterns, and labels. @@ -902,7 +905,8 @@ class Pattern(LockableImpl, AnnotatableImpl, metaclass=AutoSlots): offset: vector2 = (0., 0.), line_color: str = 'k', fill_color: str = 'none', - overdraw: bool = False): + overdraw: bool = False, + ) -> None: """ Draw a picture of the Pattern and wait for the user to inspect it diff --git a/masque/subpattern.py b/masque/subpattern.py index 8eebc95..89e143e 100644 --- a/masque/subpattern.py +++ b/masque/subpattern.py @@ -4,7 +4,7 @@ """ #TODO more top-level documentation -from typing import Dict, Tuple, Optional, Sequence, TYPE_CHECKING, Any +from typing import Dict, Tuple, Optional, Sequence, TYPE_CHECKING, Any, TypeVar import copy import numpy # type: ignore @@ -22,6 +22,9 @@ if TYPE_CHECKING: from . import Pattern +S = TypeVar('S', bound='SubPattern') + + class SubPattern(PositionableImpl, DoseableImpl, RotatableImpl, ScalableImpl, Mirrorable, PivotableImpl, Copyable, RepeatableImpl, LockableImpl, AnnotatableImpl, metaclass=AutoSlots): @@ -55,7 +58,7 @@ class SubPattern(PositionableImpl, DoseableImpl, RotatableImpl, ScalableImpl, Mi annotations: Optional[annotations_t] = None, locked: bool = False, identifier: Tuple[Any, ...] = (), - ): + ) -> None: """ Args: pattern: Pattern to reference. @@ -150,13 +153,13 @@ class SubPattern(PositionableImpl, DoseableImpl, RotatableImpl, ScalableImpl, Mi return pattern - def rotate(self, rotation: float) -> 'SubPattern': + def rotate(self: S, rotation: float) -> S: self.rotation += rotation if self.repetition is not None: self.repetition.rotate(rotation) return self - def mirror(self, axis: int) -> 'SubPattern': + def mirror(self: S, axis: int) -> S: self.mirrored[axis] = not self.mirrored[axis] self.rotation *= -1 if self.repetition is not None: @@ -176,7 +179,7 @@ class SubPattern(PositionableImpl, DoseableImpl, RotatableImpl, ScalableImpl, Mi return None return self.as_pattern().get_bounds() - def lock(self) -> 'SubPattern': + def lock(self: S) -> S: """ Lock the SubPattern, disallowing changes @@ -188,7 +191,7 @@ class SubPattern(PositionableImpl, DoseableImpl, RotatableImpl, ScalableImpl, Mi LockableImpl.lock(self) return self - def unlock(self) -> 'SubPattern': + def unlock(self: S) -> S: """ Unlock the SubPattern @@ -200,7 +203,7 @@ class SubPattern(PositionableImpl, DoseableImpl, RotatableImpl, ScalableImpl, Mi self.mirrored.flags.writeable = True return self - def deeplock(self) -> 'SubPattern': + def deeplock(self: S) -> S: """ Recursively lock the SubPattern and its contained pattern @@ -212,7 +215,7 @@ class SubPattern(PositionableImpl, DoseableImpl, RotatableImpl, ScalableImpl, Mi self.pattern.deeplock() return self - def deepunlock(self) -> 'SubPattern': + def deepunlock(self: S) -> S: """ Recursively unlock the SubPattern and its contained pattern