From e3b5492dfd420371890a301823c3896d80fc7249 Mon Sep 17 00:00:00 2001 From: Eivind Fonn Date: Thu, 15 Feb 2024 13:37:14 +0100 Subject: [PATCH] Mypy type safety: round 8 --- splipy/io/spl.py | 43 +++++++------ splipy/io/stl.py | 162 ++++++++++++++++++++++++++++++----------------- 2 files changed, 127 insertions(+), 78 deletions(-) diff --git a/splipy/io/spl.py b/splipy/io/spl.py index 938cc8e..5d0600e 100644 --- a/splipy/io/spl.py +++ b/splipy/io/spl.py @@ -1,6 +1,10 @@ from itertools import islice +from pathlib import Path +from typing import Union, TextIO, Type, Optional, Iterator +from types import TracebackType import numpy as np +from typing_extensions import Self from ..curve import Curve from ..surface import Surface @@ -13,21 +17,29 @@ class SPL(MasterIO): - def __init__(self, filename): - if not filename.endswith('.spl'): - filename += '.spl' - self.filename = filename - self.trimming_curves = [] + filename: str + fstream: TextIO - def __enter__(self): - self.fstream = open(self.filename, 'r') + def __init__(self, filename: Union[Path, str]) -> None: + self.filename = str(filename) + + def __enter__(self) -> Self: + self.fstream = open(self.filename, 'r').__enter__() return self - def lines(self): + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType] + ) -> None: + self.fstream.__exit__(exc_type, exc_val, exc_tb) + + def lines(self) -> Iterator[str]: for line in self.fstream: yield line.split('#', maxsplit=1)[0].strip() - def read(self): + def read(self) -> list[SplineObject]: lines = self.lines() version = next(lines).split() @@ -46,19 +58,12 @@ def read(self): knots = [[float(k) for k in islice(lines, nkts)] for nkts in nknots] bases = [BSplineBasis(p, kts, -1) for p, kts in zip(orders, knots)] - cpts = np.array([float(k) for k in islice(lines, totcoeffs * physdim)]) + cpts = np.array([float(k) for k in islice(lines, totcoeffs * physdim)], dtype=float) cpts = cpts.reshape(physdim, *(ncoeffs[::-1])).transpose() - if pardim == 1: - patch = Curve(*bases, controlpoints=cpts, raw=True) - elif pardim == 2: - patch = Surface(*bases, controlpoints=cpts, raw=True) - elif pardim == 3: - patch = Volume(*bases, controlpoints=cpts, raw=True) + if 1 <= pardim <= 3: + patch = SplineObject.constructor(pardim)(bases, cpts, raw=True) else: patch = SplineObject(bases, controlpoints=cpts, raw=True) return [patch] - - def __exit__(self, exc_type, exc_value, traceback): - self.fstream.close() diff --git a/splipy/io/stl.py b/splipy/io/stl.py index 7459106..f70ae8f 100644 --- a/splipy/io/stl.py +++ b/splipy/io/stl.py @@ -1,13 +1,18 @@ -#coding:utf-8 - import struct +from abc import ABC, abstractmethod +from typing import Union, Optional, Type, Sequence, TextIO, BinaryIO, cast +from types import TracebackType +from pathlib import Path import numpy as np +from typing_extensions import Self from ..surface import Surface from ..volume import Volume from ..utils import ensure_listlike from ..splinemodel import SplineModel +from ..splineobject import SplineObject +from ..types import Scalars from .master import MasterIO @@ -25,28 +30,29 @@ BINARY_FACET = "12fH" -class ASCII_STL_Writer(object): - """ Export 3D objects build of 3 or 4 vertices as ASCII STL file. - """ - def __init__(self, stream): - self.fp = stream - self._write_header() +Face = Sequence[Scalars] - def _write_header(self): - self.fp.write("solid python\n") - def close(self): - self.fp.write("endsolid python\n") +class StlWriter(ABC): - def _write(self, face): - self.fp.write(ASCII_FACET.format(face=face)) + @abstractmethod + def _write_header(self) -> None: + ... + + @abstractmethod + def _write(self, face: Face) -> None: + ... - def _split(self, face): + @abstractmethod + def close(self) -> None: + ... + + def _split(self, face: Face) -> tuple[Face, Face]: p1, p2, p3, p4 = face return (p1, p2, p3), (p3, p4, p1) - def add_face(self, face): - """ Add one face with 3 or 4 vertices. """ + def add_face(self, face: Face) -> None: + """Add one face with 3 or 4 vertices.""" if len(face) == 4: face1, face2 = self._split(face) self._write(face1) @@ -56,30 +62,50 @@ def add_face(self, face): else: raise ValueError('only 3 or 4 vertices for each face') - def add_faces(self, faces): + def add_faces(self, faces: Sequence[Face]) -> None: """ Add many faces. """ for face in faces: self.add_face(face) -class BINARY_STL_Writer(ASCII_STL_Writer): - """ Export 3D objects build of 3 or 4 vertices as binary STL file. - """ - def __init__(self, stream): - self.counter = 0 - #### new-style classes way of calling super constructor - # super(Binary_STL_Writer, self).__init__(stream) - #### old-style classes way of doing it - ASCII_STL_Writer.__init__(self, stream) +class AsciiStlWriter(StlWriter): + """Export 3D objects build of 3 or 4 vertices as ASCII STL file.""" - def close(self): + fp: TextIO + + def __init__(self, stream: TextIO) -> None: + self.fp = stream self._write_header() - def _write_header(self): + def _write_header(self) -> None: + self.fp.write("solid python\n") + + def close(self) -> None: + self.fp.write("endsolid python\n") + + def _write(self, face: Face) -> None: + self.fp.write(ASCII_FACET.format(face=face)) + + +class BinaryStlWriter(StlWriter): + """Export 3D objects build of 3 or 4 vertices as binary STL file.""" + + counter: int + fp: BinaryIO + + def __init__(self, stream: BinaryIO) -> None: + self.counter = 0 + self.fp = stream + self._write_header() + + def close(self) -> None: + self._write_header() + + def _write_header(self) -> None: self.fp.seek(0) self.fp.write(struct.pack(BINARY_HEADER, b'Python Binary STL Writer', self.counter)) - def _write(self, face): + def _write(self, face: Face) -> None: self.counter += 1 data = [ 0., 0., 0., @@ -92,42 +118,65 @@ def _write(self, face): class STL(MasterIO): - def __init__(self, filename, binary=True): - if filename[-4:] != '.stl': - filename += '.stl' - self.filename = filename - self.binary = binary - def __enter__(self): + filename: str + binary: bool + + writer: Union[BinaryStlWriter, AsciiStlWriter] + + def __init__(self, filename: Union[str, Path], binary: bool = True) -> None: + self.filename = str(filename) + self.binary = binary + + def __enter__(self) -> Self: if self.binary: - fp = open(self.filename, 'wb') - self.writer = BINARY_STL_Writer(fp) + self.writer = BinaryStlWriter(open(self.filename, 'wb')) else: - fp = open(self.filename, 'w') - self.writer = ASCII_STL_Writer(fp) + self.writer = AsciiStlWriter(open(self.filename, 'w')) return self - def write(self, obj, n=None): + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType] + ) -> None: + self.writer.close() + self.writer.fp.__exit__(exc_type, exc_val, exc_tb) + + def write( + self, + obj: Union[SplineModel, Sequence[SplineObject], SplineObject], + n: Optional[Union[int, Sequence[int]]] = None, + ) -> None: if isinstance(obj, SplineModel): if obj.pardim == 3: # volume model - for surface in obj.boundary(): - self.write_surface(surface.obj,n) + for node in obj.boundary(): + self.write_surface(cast(Surface, node.obj), n) elif obj.pardim == 2: # surface model for surface in obj: - self.write_surface(surface, n) + self.write_surface(cast(Surface, surface), n) elif isinstance(obj, Volume): - for surface in obj.faces(): - if surface is not None: # happens with periodic volumes - self.write_surface(surface, n) + for face in obj.faces(): + if face is not None: # happens with periodic volumes + self.write_surface(face, n) elif isinstance(obj, Surface): self.write_surface(obj, n) + elif isinstance(obj, Sequence): + for o in obj: + self.write(o) + else: raise ValueError('Unsopported object for STL format') - def write_surface(self, surface, n=None): + def write_surface( + self, + surface: Surface, + n: Optional[Union[int, Sequence[int]]] = None, + ) -> None: # choose evaluation points as one of three cases: # 1. specified with input # 2. linear splines, only picks knots @@ -142,9 +191,9 @@ def write_surface(self, surface, n=None): else: knots = surface.knots(0) p = surface.order(0) - u = [np.linspace(k0,k1, 2*p-3, endpoint=False) for (k0,k1) in zip(knots[:-1], knots[1:])] - u = [point for element in u for point in element] + list(knots) - u = np.sort(u) + ut = [np.linspace(k0,k1, 2*p-3, endpoint=False) for (k0,k1) in zip(knots[:-1], knots[1:])] + ut = [point for element in ut for point in element] + list(knots) + u = np.sort(ut) if n is not None: v = np.linspace(surface.start(1), surface.end(1), n[1]) @@ -153,9 +202,9 @@ def write_surface(self, surface, n=None): else: knots = surface.knots(1) p = surface.order(1) - v = [np.linspace(k0,k1, 2*p-3, endpoint=False) for (k0,k1) in zip(knots[:-1], knots[1:])] - v = [point for element in v for point in element] + list(knots) - v = np.sort(v) + vt = [np.linspace(k0,k1, 2*p-3, endpoint=False) for (k0,k1) in zip(knots[:-1], knots[1:])] + vt = [point for element in vt for point in element] + list(knots) + v = np.sort(vt) # perform evaluation and make sure that we have 3 components (in case of 2D geometries) x = surface(u,v) @@ -166,8 +215,3 @@ def write_surface(self, surface, n=None): faces = [[x[i,j], x[i,j+1], x[i+1,j+1], x[i+1,j]] for i in range(x.shape[0]-1) for j in range(x.shape[1]-1)] self.writer.add_faces(faces) - - def __exit__(self, exc_type, exc_value, traceback): - self.writer.close() - self.writer.fp.close() -