Mypy type safety: round 9

splipy/io/grdecl.py

@@ -1,6 +1,7 @@
 from itertools import product, chain
 import re
 import warnings
+from typing import Union, Sequence
 
 import numpy as np
 from tqdm import tqdm
@@ -12,6 +13,8 @@
 from ..volume import Volume
 from ..basis import BSplineBasis
 from ..utils import ensure_listlike
+from ..splineobject import SplineObject
+from ..splinemodel import SplineModel
 from .. import surface_factory, curve_factory, volume_factory
 
 from .master import MasterIO
@@ -211,6 +214,8 @@ def read(self):
 
         self.raw = DiscontBoxMesh(self.n, self.coord, self.zcorn)
 
+    def write(self, obj: Union[SplineObject, SplineModel, Sequence[SplineObject]]) -> None:
+        raise IOError('Writing to GRDECL not supported')
 
     def get_c0_mesh(self):
         # Create the C0-mesh

splipy/io/master.py

@@ -6,9 +6,10 @@
 from typing_extensions import Self
 
 from ..splineobject import SplineObject
+from ..splinemodel import SplineModel
 
 
-class MasterIO:
+class MasterIO(ABC):
 
     def __init__(self, filename: Union[str, Path]) -> None:
         """Create an IO object attached to a file.
@@ -31,7 +32,7 @@ def __exit__(
         ...
 
     @abstractmethod
-    def write(self, obj: Union[SplineObject, Sequence[SplineObject]]) -> None:
+    def write(self, obj: Union[SplineObject, Sequence[SplineObject], SplineModel]) -> None:
         """Write one or more objects to the file.
 
         :param obj: The object(s) to write

splipy/io/ofoam.py

@@ -1,43 +1,56 @@
+from pathlib import Path
 from itertools import groupby
 from operator import itemgetter
 from os.path import exists, isdir, join
 from os import makedirs
+from typing import Union, Optional, Type, Sequence
+from types import TracebackType
+from typing_extensions import Self
 
 import numpy as np
 
+from .master import MasterIO
+from ..splineobject import SplineObject
 from ..splinemodel import SplineModel
 
 
-class OpenFOAM(object):
+class OpenFOAM(MasterIO):
+    target: str
 
-    def __init__(self, target):
-        self.target = target
+    def __init__(self, target: Union[Path, str]) -> None:
+        self.target = str(target)
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         # Create the target directory if it does not exist
         if not exists(self.target):
             makedirs(self.target)
+
         # If it does, ensure that it's a directory
         elif not isdir(self.target):
             raise FileExistsError('{} exists and is not a directory'.format(self.target))
 
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
-        pass
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType]
+    ) -> None:
+        return
 
-    def _header(self, cls, obj, note=None):
+    def _header(self, cls: str, obj: str, note: Optional[str] = None) -> str:
         s = 'FoamFile\n{\n'
         s += '    version     2.0;\n'
         s += '    format      ascii;\n'
-        s += '    class       %s;\n' % cls
+        s += f'    class       {cls};\n'
         if note:
-            s += '    note        "%s";\n' % note
-        s += '    object      %s;\n' % obj
+            s += f'    note        "{note}";\n'
+        s += f'    object      {obj};\n'
         s += '}\n'
         return s
 
-    def write(self, model):
+    def write(self, model: Union[SplineObject, Sequence[SplineObject], SplineModel]) -> None:
         assert isinstance(model, SplineModel), "OpenFOAM.write only supports SplineModel objects"
 
         # Only linear volumes in 3D, please
@@ -60,11 +73,11 @@ def write(self, model):
         # - All faces in the same boundary must be contiguous
         # - Low number owners before high number owners
         # - Low number neighbors before high number neighbors
-        faces = list(faces)
-        faces = sorted(faces, key=itemgetter('neighbor'))
-        faces = sorted(faces, key=itemgetter('owner'))
-        faces = sorted(faces, key=lambda x: (x['name'] is not None, x['name']))
-        faces = np.array(faces)
+        faces_list = list(faces)
+        faces_list = sorted(faces_list, key=itemgetter('neighbor'))
+        faces_list = sorted(faces_list, key=itemgetter('owner'))
+        faces_list = sorted(faces_list, key=lambda x: (x['name'] is not None, x['name']))
+        faces = np.array(faces_list)
 
         # Write the points file (vertex coordinates)
         with open(join(self.target, 'points'), 'w') as f:

splipy/io/spl.py

@@ -1,6 +1,6 @@
 from itertools import islice
 from pathlib import Path
-from typing import Union, TextIO, Type, Optional, Iterator
+from typing import Union, TextIO, Type, Optional, Iterator, Sequence
 from types import TracebackType
 
 import numpy as np
@@ -10,6 +10,7 @@
 from ..surface import Surface
 from ..volume import Volume
 from ..splineobject import SplineObject
+from ..splinemodel import SplineModel
 from ..basis import BSplineBasis
 
 from .master import MasterIO
@@ -39,6 +40,9 @@ def lines(self) -> Iterator[str]:
         for line in self.fstream:
             yield line.split('#', maxsplit=1)[0].strip()
 
+    def write(self, obj: Union[SplineObject, SplineModel, Sequence[SplineObject]]) -> None:
+        raise IOError('Writing to SPL not supported')
+
     def read(self) -> list[SplineObject]:
         lines = self.lines()
 

splipy/io/stl.py

@@ -144,6 +144,9 @@ def __exit__(
         self.writer.close()
         self.writer.fp.__exit__(exc_type, exc_val, exc_tb)
 
+    def read(self) -> list[SplineObject]:
+        raise IOError('Reading STL not supported')
+
     def write(
         self,
         obj: Union[SplineModel, Sequence[SplineObject], SplineObject],

splipy/io/svg.py

@@ -11,6 +11,7 @@
 from ..curve import Curve
 from ..surface import Surface
 from ..splineobject import SplineObject
+from ..splinemodel import SplineModel
 from ..basis import BSplineBasis
 from ..types import FArray
 from .. import curve_factory, state
@@ -214,7 +215,7 @@ def write_surface(self, surface: Surface, fill: str = '#ffcc99') -> None:
         for meshline in knotlines:
             self.write_curve(groupNode, meshline, width=1)
 
-    def write(self, obj: Union[SplineObject, Sequence[SplineObject]]) -> None:
+    def write(self, obj: Union[SplineObject, SplineModel, Sequence[SplineObject]]) -> None:
         """Writes a list of planar curves and surfaces to vector graphics SVG file.
         The image will never be stretched, and the geometry will keep width/height ratio
         of original geometry, regardless of provided width/height ratio from arguments.

splipy/io/threedm.py

@@ -8,6 +8,7 @@
 import numpy as np
 import rhino3dm as rhino
 
+from ..splinemodel import SplineModel
 from ..splineobject import SplineObject
 from ..curve import Curve
 from ..surface import Surface
@@ -61,7 +62,7 @@ def __exit__(
     ) -> None:
         pass
 
-    def write(self, obj: Union[SplineObject, Sequence[SplineObject]]) -> None:
+    def write(self, obj: Union[SplineObject, SplineModel, Sequence[SplineObject]]) -> None:
         raise IOError('Writing to 3DM not supported')
 
     def read(self) -> list[SplineObject]: