ENH: create FileWrapper method for retrieving text stream (#1492)

Author: ErichSuter

src/xtgeo/io/_file.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import io
 import os
 import pathlib
@@ -12,7 +13,9 @@
 from enum import Enum
 from os.path import join
 from tempfile import mkstemp
-from typing import TYPE_CHECKING, Literal, Union
+from typing import TYPE_CHECKING, Generator, Literal, TextIO, Union
+
+from typing_extensions import Self
 
 import xtgeo._cxtgeo
 from xtgeo.common.exceptions import InvalidFileFormatError
@@ -646,3 +649,28 @@ def _format_from_contents(self) -> FileFormat:
                 return FileFormat.TSURF
 
         return FileFormat.UNKNOWN
+
+    @contextlib.contextmanager
+    def get_text_stream(self: Self) -> Generator[TextIO, None, None]:
+        """
+        Context manager to handle both file paths and file-like objects for reading.
+        Yields:
+            A text stream (TextIO) for reading.
+        Raises:
+            FileNotFoundError: If the file does not exist.
+        """
+
+        if not self.check_file():
+            raise FileNotFoundError(f"\nFile {self.name}:\nThe file does not exist.")
+
+        if isinstance(self.file, pathlib.Path):
+            with open(self.file, "r") as stream:
+                yield stream
+        elif isinstance(self.file, io.BytesIO):
+            with io.TextIOWrapper(self.file) as text_wrapper:
+                text_wrapper.seek(0)
+                yield text_wrapper
+        else:
+            # StringIO is already a text stream
+            self.file.seek(0)
+            yield self.file

src/xtgeo/io/tsurf/_tsurf_reader.py

@@ -1,8 +1,5 @@
 import warnings
-from contextlib import contextmanager
 from dataclasses import dataclass
-from io import BytesIO, TextIOWrapper
-from pathlib import Path
 from typing import Any, Generator, TextIO
 
 import numpy as np
@@ -681,34 +678,8 @@ def _parse_tsurf(stream: TextIO, filepath_errmsg: str) -> TSurfData:
     return tsurf_data
 
 
-@contextmanager
-def _get_text_stream(
-    wrapped_file: FileWrapper,
-    encoding: str = "utf-8",
-) -> Generator[TextIO, None, None]:
-    """Context manager that yields a text stream."""
-
-    if not wrapped_file or not wrapped_file.check_file():
-        raise FileNotFoundError(
-            f"\nIn file {wrapped_file.name}:\nThe file does not exist."
-        )
-
-    wrapped_file.fileformat(FileFormat.TSURF.value[0], strict=True)
-
-    if isinstance(wrapped_file.file, Path):
-        with open(wrapped_file.file, encoding=encoding) as stream:
-            yield stream
-    elif isinstance(wrapped_file.file, BytesIO):
-        with TextIOWrapper(wrapped_file.file, encoding=encoding) as text_wrapper:
-            yield text_wrapper
-    else:
-        # StringIO is already a text stream
-        yield wrapped_file.file
-
-
 def read_tsurf(
     file: FileLike,
-    encoding: str = "utf-8",
 ) -> TSurfData:
     """
     Read a file on the TSURF format and parse its triangulated surface data.
@@ -742,5 +713,11 @@ def read_tsurf(
 
     wrapped_file = FileWrapper(file)
 
-    with _get_text_stream(wrapped_file, encoding) as stream:
-        return _parse_tsurf(stream, str(wrapped_file.name))
+    if not wrapped_file.check_file():
+        raise FileNotFoundError(
+            f"\nIn file {wrapped_file.name}:\nThe file does not exist."
+        )
+    wrapped_file.fileformat(FileFormat.TSURF.value[0], strict=True)
+
+    with wrapped_file.get_text_stream() as stream:
+        return _parse_tsurf(stream, filepath_errmsg=str(wrapped_file.name))

src/xtgeo/xyz/_xyz_io.py

@@ -4,9 +4,8 @@
 
 import contextlib
 from copy import deepcopy
-from io import BytesIO, StringIO, TextIOWrapper
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Generator, Iterable, Literal, TextIO, TypeGuard
+from io import StringIO
+from typing import TYPE_CHECKING, Any, Iterable, Literal, TypeGuard
 
 import numpy as np
 import pandas as pd
@@ -252,34 +251,6 @@ def import_zmap(
     return {"xname": xname, "yname": yname, "zname": zname, "values": df}
 
 
-@contextlib.contextmanager
-def _get_text_stream(
-    wrapped_file: FileWrapper, encoding: str = "utf-8"
-) -> Generator[TextIO, None, None]:
-    """Context manager to handle both file paths and file-like objects for reading."""
-
-    # TODO: consider moving method to xtgeo.io._file: FileWrapper
-    # Ensure that the method appropriately handles the different
-    # configurations/file types of the FileWrapper.
-
-    if not wrapped_file or not wrapped_file.check_file():
-        raise FileNotFoundError(
-            f"\nFile {wrapped_file.name}:\nThe file does not exist."
-        )
-
-    if isinstance(wrapped_file.file, Path):
-        with open(wrapped_file.file, "r", encoding=encoding) as stream:
-            yield stream
-    elif isinstance(wrapped_file.file, BytesIO):
-        with TextIOWrapper(wrapped_file.file, encoding=encoding) as text_wrapper:
-            text_wrapper.seek(0)
-            yield text_wrapper
-    else:
-        # StringIO is already a text stream
-        wrapped_file.file.seek(0)
-        yield wrapped_file.file
-
-
 def import_rms_attr(pfile: FileWrapper, zname: str = "Z_TVDSS") -> dict[str, Any]:
     """The RMS ascii file Points format with attributes.
 
@@ -323,7 +294,7 @@ def import_rms_attr(pfile: FileWrapper, zname: str = "Z_TVDSS") -> dict[str, Any
     # parse header
     skiprows = 0
 
-    with _get_text_stream(pfile) as rmsfile:
+    with pfile.get_text_stream() as rmsfile:
         for iline in range(20):
             fields = rmsfile.readline().split()
             if len(fields) != 2:

tests/test_io/test_file.py

@@ -253,6 +253,59 @@ def test_fileformat_provided_prefer_given(testdata_path, filename, expected_form
         xtgeo_file.fileformat(fileformat="segy", strict=True) == FileFormat.SEGY
 
 
+def test_get_text_stream(testdata_path: str) -> None:
+    """Test getting a text stream from a file."""
+    file_path = pathlib.Path(testdata_path) / "surfaces/drogon/3/F5.ts"
+    xtgeo_file = FileWrapper(file_path)
+    with xtgeo_file.get_text_stream() as stream:
+        assert isinstance(stream, io.TextIOWrapper)
+        assert stream.readline().startswith("GOCAD TSurf 1")
+
+
+@pytest.mark.parametrize("filename", ["NOSUCH.EGRID", "NOSUCH/NOSUCH.EGRID"])
+def test_get_text_stream_file_does_not_exist(
+    reek_grid_path: pathlib.Path, filename: str
+) -> None:
+    """Test getting a text stream from a non-existing file."""
+    xtgeo_file = FileWrapper(reek_grid_path / filename)
+    with pytest.raises(FileNotFoundError, match="file does not exist"):
+        # Use __enter__ (not 'pass') to ensure proper test coverage
+        xtgeo_file.get_text_stream().__enter__()
+
+
+def test_get_text_stream_from_binary_file(testdata_path: str) -> None:
+    """Test getting a text stream from a binary file."""
+    file_path = pathlib.Path(testdata_path) / "3dgrids/reek/REEK.EGRID"
+    xtgeo_file = FileWrapper(file_path)
+
+    # Document expected behavior:
+    # UnicodeDecodeError raised when trying to read binary file as text
+    with pytest.raises(UnicodeDecodeError), xtgeo_file.get_text_stream() as stream:
+        stream.readlines()
+
+
+def test_get_text_stream_empty_file(testdata_path: str) -> None:
+    """Test getting a text stream from an empty text file."""
+    empty_file_path = pathlib.Path(testdata_path) / "empty_file.txt"
+    empty_file_path.touch()
+    xtgeo_file = FileWrapper(empty_file_path)
+    with xtgeo_file.get_text_stream() as stream:
+        assert len(stream.readlines()) == 0
+
+
+def test_get_text_stream_from_memstream() -> None:
+    """Test getting a text stream from an in-memory stream."""
+    dummy_text = "Line 1\nLine 2\nLine 3\n"
+    with FileWrapper(io.StringIO()).get_text_stream() as stream:
+        assert len(stream.readlines()) == 0
+    with FileWrapper(io.BytesIO()).get_text_stream() as stream:
+        assert len(stream.readlines()) == 0
+    with FileWrapper(io.StringIO(dummy_text)).get_text_stream() as stream:
+        assert len(stream.readlines()) == 3
+    with FileWrapper(io.BytesIO(dummy_text.encode())).get_text_stream() as stream:
+        assert len(stream.readlines()) == 3
+
+
 @pytest.mark.parametrize("strict", [False, True])
 @pytest.mark.parametrize("filename", SURFACE_FILE_FORMATS.keys())
 def test_fileformat_hdf_stream(testdata_path, filename, strict):

tests/test_io/test_tsurf/test_tsurf_reader.py

@@ -1,4 +1,4 @@
-from io import BytesIO, StringIO
+from io import StringIO
 from pathlib import Path
 
 import pytest
@@ -115,55 +115,6 @@ def test_file_string_input(tmp_path: str, complete_tsurf_file) -> None:
     assert result_path is not None
 
 
-def test_file_path_input(tmp_path: Path, complete_tsurf_file) -> None:
-    """Test reading from Path input."""
-    # Test with Path
-    filepath = tmp_path / "test.ts"
-    with open(filepath, "w") as f:
-        f.write(complete_tsurf_file)
-
-    result_path = read_tsurf(filepath)
-    assert result_path is not None
-
-
-def test_file_stringio_input(complete_tsurf_file) -> None:
-    """Test reading from StringIO input."""
-    result_stringio = read_tsurf(tsurf_stream(complete_tsurf_file))
-    assert result_stringio is not None
-
-
-def test_file_bytesio_input(complete_tsurf_file) -> None:
-    """Test reading from BytesIO input."""
-    result_bytesio = read_tsurf(BytesIO(complete_tsurf_file.encode("utf-8")))
-    assert result_bytesio is not None
-
-
-def test_file_non_regular_file_input(tmp_path: Path) -> None:
-    """Test reading from a non-regular file (e.g., folder)."""
-
-    non_regular_file = tmp_path / "some_folder"
-    non_regular_file.mkdir()
-
-    with pytest.raises(FileNotFoundError, match="does not exist"):
-        read_tsurf(non_regular_file)
-
-
-def test_file_other_than_filelike_input() -> None:
-    """Test reading from an unsupported input type."""
-    with pytest.raises(
-        RuntimeError, match="Cannot instantiate <class 'xtgeo.io._file.FileWrapper'>"
-    ):
-        read_tsurf(12345)  # Invalid input type
-
-
-def test_file_empty():
-    """Test that empty file raises error."""
-    with pytest.raises(
-        ValueError, match="does not match format detected from file contents"
-    ):
-        read_tsurf(StringIO(""))
-
-
 def test_file_unusual_suffix(minimal_tsurf_file, tmp_path: Path) -> None:
     """
     Test with unusual file suffix.
@@ -176,13 +127,6 @@ def test_file_unusual_suffix(minimal_tsurf_file, tmp_path: Path) -> None:
     assert result_unusual_suffix is not None
 
 
-def test_file_non_existent(tmp_path: Path) -> None:
-    """Test with non-existent file."""
-    filepath = tmp_path / "non_existent.ts"
-    with pytest.raises(FileNotFoundError, match="does not exist"):
-        read_tsurf(filepath)
-
-
 def test_comments_and_empty_lines(tmp_path: Path) -> None:
     """Test handling of comments and empty lines in files."""
     content_lines = [