ENH: add tsurf writer

Author: ErichSuter

src/xtgeo/io/_file.py

@@ -13,7 +13,13 @@
 from enum import Enum
 from os.path import join
 from tempfile import mkstemp
-from typing import TYPE_CHECKING, Generator, Literal, TextIO, Union
+from typing import (
+    TYPE_CHECKING,
+    Generator,
+    Literal,
+    TextIO,
+    Union,
+)
 
 from typing_extensions import Self
 
@@ -674,3 +680,43 @@ def get_text_stream(self: Self) -> Generator[TextIO, None, None]:
             # StringIO is already a text stream
             self.file.seek(0)
             yield self.file
+
+    @contextlib.contextmanager
+    def get_text_stream_write(self: Self) -> Generator[TextIO, None, None]:
+        """
+        Context manager to handle both file paths and file-like objects for writing.
+
+        Args:
+            encoding: Character encoding to use for text files. Defaults to "utf-8".
+
+        Yields:
+            A text stream (TextIO) for writing.
+
+        Raises:
+            OSError: If the parent folder does not exist or is not writable.
+
+        Example::
+            >>> wrapper = FileWrapper("output.txt", mode="w")
+            >>> with wrapper.get_text_stream_for_writing() as stream:
+            ...     stream.write("Hello, world!\\n")
+        """
+
+        encoding: str = "utf-8"
+
+        if isinstance(self.file, pathlib.Path):
+            self.check_folder(raiseerror=OSError)
+            with open(self.file, "w", encoding=encoding) as stream:
+                yield stream
+        elif isinstance(self.file, io.BytesIO):
+            # Wrap BytesIO in TextIOWrapper for text writing
+            text_wrapper = io.TextIOWrapper(
+                self.file, encoding=encoding, write_through=True
+            )
+            try:
+                yield text_wrapper
+            finally:
+                # Detach the wrapper to prevent it from closing the underlying BytesIO
+                text_wrapper.detach()
+        else:
+            # StringIO is already a text stream, just yield it
+            yield self.file

src/xtgeo/io/tsurf/_tsurf_reader.py

@@ -1,3 +1,5 @@
+from io import BytesIO, StringIO
+import pathlib
 import warnings
 from dataclasses import dataclass
 from typing import Any, Generator, TextIO
@@ -790,4 +792,61 @@ def to_file(
             FileNotFoundError: If file path doesn't exist or isn't a regular file
         """
 
-        raise NotImplementedError("TSurfData.to_file() is not yet implemented.")
+        wrapped_file = FileWrapper(file)
+        wrapped_file.check_folder(raiseerror=OSError)
+
+        encoding = "utf-8"
+
+        # TODO: remove commented code below
+        # Looked in polygons.py, only need to open the file and write to it
+        # But this is when wrapped_file is a str
+        # with open(wrapped_file, "w") as fout:
+        #     for col in use_attributes:
+        #         if col in df.columns:
+        #             fout.write(transl[xyz._attrs[col]] + " " + col + "\n")
+
+        lines = []
+
+        # TSurf signature line
+        lines.append("GOCAD TSurf 1\n")
+
+        lines.append("HEADER {\n")
+        lines.append(f"name: {self.header.name}\n")
+        lines.append("}\n")
+
+        # Optional: coordinate system
+        if self.coord_sys:
+            lines.append("GOCAD_ORIGINAL_COORDINATE_SYSTEM\n")
+            lines.append(f"NAME {self.coord_sys.name}\n")
+
+            axis_name_str = " ".join([f'"{name}"' for name in self.coord_sys.axis_name])
+            lines.append(f"AXIS_NAME {axis_name_str}\n")
+
+            axis_unit_str = " ".join([f'"{unit}"' for unit in self.coord_sys.axis_unit])
+            lines.append(f"AXIS_UNIT {axis_unit_str}\n")
+
+            lines.append(f"ZPOSITIVE {self.coord_sys.zpositive}\n")
+            lines.append("END_ORIGINAL_COORDINATE_SYSTEM\n")
+
+        # TFACE section with vertices and triangles
+        lines.append("TFACE\n")
+        for i, vertex in enumerate(self.vertices, start=1):
+            lines.append(f"VRTX {i} {vertex[0]} {vertex[1]} {vertex[2]} CNXYZ\n")
+        for triangle in self.triangles:
+            lines.append(f"TRGL {triangle[0]} {triangle[1]} {triangle[2]}\n")
+        lines.append("END\n")
+
+        with wrapped_file.get_text_stream_write() as stream:
+            stream.writelines(lines)
+
+        # TODO: remove this
+        # Alternative way of writing:
+        # if isinstance(wrapped_file.file, pathlib.Path):
+        #     with open(str(wrapped_file.name), "w", encoding=encoding) as f_out:
+        #         f_out.writelines(lines)
+        # else:
+        #     wrapped_file.file.seek(0)
+        #     if isinstance(wrapped_file.file, BytesIO):
+        #         wrapped_file.file.write("".join(lines).encode(encoding))
+        #     else:  # isinstance(wrapped_file.file, StringIO):
+        #         wrapped_file.file.write("".join(lines))

tests/test_io/test_tsurf/test_tsurf_reader.py

@@ -1,5 +1,5 @@
 from dataclasses import FrozenInstanceError
-from io import StringIO
+from io import BytesIO, StringIO
 from pathlib import Path
 
 import numpy as np
@@ -113,8 +113,8 @@ def test_file_string_input(tmp_path: str, complete_tsurf_file: str) -> None:
     with open(filepath, "w") as f:
         f.write(complete_tsurf_file)
 
-    result_path = TSurfData.from_file(filepath)
-    assert result_path is not None
+    result = TSurfData.from_file(filepath)
+    assert result is not None
 
 
 def test_file_unusual_suffix(minimal_tsurf_file: str, tmp_path: Path) -> None:
@@ -1041,3 +1041,64 @@ def test_tsurfdata_get_vertices_return_same_reference() -> None:
     )
 
     assert data.get_vertices is data.vertices
+
+
+def test_tsurf_data_roundtrip_write_to_file(
+    complete_tsurf_file: str, tmp_path: Path
+) -> None:
+    """Test writing TSurfData."""
+    result = TSurfData.from_file(tsurf_stream(complete_tsurf_file))
+    assert result is not None
+
+    output_filepath = tmp_path / "output.ts"
+    result.to_file(output_filepath)
+
+    result_written = TSurfData.from_file(output_filepath)
+    assert result_written is not None
+
+    assert result.header == result_written.header
+    assert result.coord_sys == result_written.coord_sys
+    assert np.array_equal(result.vertices, result_written.vertices)
+    assert np.array_equal(result.triangles, result_written.triangles)
+
+
+def test_tsurf_data_roundtrip_string_io(complete_tsurf_file: str) -> None:
+    """Test writing and reading TSurfData using StringIO."""
+    result = TSurfData.from_file(tsurf_stream(complete_tsurf_file))
+    assert result is not None
+
+    output_stream = StringIO()
+    result.to_file(output_stream)
+
+    output_stream.seek(0)
+    result_written = TSurfData.from_file(output_stream)
+    assert result_written is not None
+
+    assert result.header == result_written.header
+    assert result.coord_sys == result_written.coord_sys
+    assert np.array_equal(result.vertices, result_written.vertices)
+    assert np.array_equal(result.triangles, result_written.triangles)
+
+
+def test_tsurf_data_roundtrip_bytes_io(complete_tsurf_file: str) -> None:
+    """Test writing and reading TSurfData using BytesIO."""
+    result = TSurfData.from_file(tsurf_stream(complete_tsurf_file))
+    assert result is not None
+
+    output_stream = BytesIO()
+    result.to_file(output_stream)
+
+    output_stream.seek(0)
+    result_written = TSurfData.from_file(output_stream)
+    assert result_written is not None
+
+    assert result.header == result_written.header
+    assert result.coord_sys == result_written.coord_sys
+    assert np.array_equal(result.vertices, result_written.vertices)
+    assert np.array_equal(result.triangles, result_written.triangles)
+
+    # TODO: see dataio writer tests
+    # TODO: test with invalid filepaths/streams, non-existing folders, etc.
+    # TODO: test with a couple of different encodings
+    # TODO: write erroneous files and check errors
+    # TODO: dataio: test_tsurf_reader_invalid_lines()