Merge pull request #1743 from OceanParcels/better-prints

New and updated `reprs` for Variable, ParticleFile, Field, VectorField, and ParticleSet

Author: erikvansebille

docs/examples/parcels_tutorial.ipynb

@@ -20,7 +20,7 @@
     assert_valid_gridindexingtype,
     assert_valid_interp_method,
 )
-from parcels.tools._helpers import deprecated_made_private, timedelta_to_float
+from parcels.tools._helpers import default_repr, deprecated_made_private, field_repr, timedelta_to_float
 from parcels.tools.converters import (
     Geographic,
     GeographicPolar,
@@ -149,6 +149,8 @@ class Field:
     * `Nested Fields <../examples/tutorial_NestedFields.ipynb>`__
     """
 
+    allow_time_extrapolation: bool
+    time_periodic: TimePeriodic
     _cast_data_dtype: type[np.float32] | type[np.float64]
 
     def __init__(
@@ -321,6 +323,9 @@ def __init__(
         if len(kwargs) > 0:
             raise SyntaxError(f'Field received an unexpected keyword argument "{list(kwargs.keys())[0]}"')
 
+    def __repr__(self) -> str:
+        return field_repr(self)
+
     @property
     @deprecated_made_private  # TODO: Remove 6 months after v3.1.0
     def dataFiles(self):
@@ -1915,6 +1920,14 @@ def __init__(self, name: str, U: Field, V: Field, W: Field | None = None):
                 assert W.interp_method == "cgrid_velocity", "Interpolation methods of U and W are not the same."
                 assert self._check_grid_dimensions(U.grid, W.grid), "Dimensions of U and W are not the same."
 
+    def __repr__(self):
+        w_repr = default_repr(self.W) if self.W is not None else repr(self.W)
+        return f"""<{type(self).__name__}>
+    name: {self.name!r}
+    U: {default_repr(self.U)}
+    V: {default_repr(self.V)}
+    W: {w_repr}"""
+
     @staticmethod
     def _check_grid_dimensions(grid1, grid2):
         return (

parcels/field.py

@@ -14,7 +14,7 @@
 from parcels.grid import Grid
 from parcels.gridset import GridSet
 from parcels.particlefile import ParticleFile
-from parcels.tools._helpers import deprecated_made_private
+from parcels.tools._helpers import deprecated_made_private, fieldset_repr
 from parcels.tools.converters import TimeConverter, convert_xarray_time_units
 from parcels.tools.loggers import logger
 from parcels.tools.statuscodes import TimeExtrapolationError
@@ -56,6 +56,9 @@ def __init__(self, U: Field | NestedField | None, V: Field | NestedField | None,
         self.compute_on_defer = None
         self._add_UVfield()
 
+    def __repr__(self):
+        return fieldset_repr(self)
+
     @property
     def particlefile(self):
         return self._particlefile

parcels/fieldset.py

@@ -53,7 +53,7 @@ def __set__(self, instance, value):
             setattr(instance, f"_{self.name}", value)
 
     def __repr__(self):
-        return f"PVar<{self.name}|{self.dtype}>"
+        return f"Variable(name={self._name}, dtype={self.dtype}, initial={self.initial}, to_write={self.to_write})"
 
     def is64bit(self):
         """Check whether variable is 64-bit."""

parcels/particle.py

@@ -10,7 +10,7 @@
 
 import parcels
 from parcels._compat import MPI
-from parcels.tools._helpers import deprecated, deprecated_made_private, timedelta_to_float
+from parcels.tools._helpers import default_repr, deprecated, deprecated_made_private, timedelta_to_float
 from parcels.tools.warnings import FileWarning
 
 __all__ = ["ParticleFile"]
@@ -116,6 +116,16 @@ def __init__(self, name, particleset, outputdt=np.inf, chunks=None, create_new_z
                 fname = name if extension in [".zarr"] else f"{name}.zarr"
         self._fname = fname
 
+    def __repr__(self) -> str:
+        return (
+            f"{type(self).__name__}("
+            f"name={self.fname!r}, "
+            f"particleset={default_repr(self.particleset)}, "
+            f"outputdt={self.outputdt!r}, "
+            f"chunks={self.chunks!r}, "
+            f"create_new_zarrfile={self.create_new_zarrfile!r})"
+        )
+
     @property
     def create_new_zarrfile(self):
         return self._create_new_zarrfile

parcels/particlefile.py

@@ -27,7 +27,7 @@
 from parcels.particle import JITParticle, Variable
 from parcels.particledata import ParticleData, ParticleDataIterator
 from parcels.particlefile import ParticleFile
-from parcels.tools._helpers import deprecated, deprecated_made_private, timedelta_to_float
+from parcels.tools._helpers import deprecated, deprecated_made_private, particleset_repr, timedelta_to_float
 from parcels.tools.converters import _get_cftime_calendars, convert_to_flat_array
 from parcels.tools.global_statics import get_package_dir
 from parcels.tools.loggers import logger
@@ -112,6 +112,7 @@ def __init__(
         self.fieldset = fieldset
         self.fieldset._check_complete()
         self.time_origin = fieldset.time_origin
+        self._pclass = pclass
 
         # ==== first: create a new subclass of the pclass that includes the required variables ==== #
         # ==== see dynamic-instantiation trick here: https://www.python-course.eu/python3_classes_and_type.php ==== #
@@ -386,8 +387,12 @@ def size(self):
         # ==== to change at some point - len and size are different things ==== #
         return len(self.particledata)
 
+    @property
+    def pclass(self):
+        return self._pclass
+
     def __repr__(self):
-        return "\n".join([str(p) for p in self])
+        return particleset_repr(self)
 
     def __len__(self):
         return len(self.particledata)

parcels/particleset.py

@@ -1,12 +1,19 @@
 """Internal helpers for Parcels."""
 
+from __future__ import annotations
+
 import functools
+import textwrap
 import warnings
 from collections.abc import Callable
 from datetime import timedelta
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
+if TYPE_CHECKING:
+    from parcels import Field, FieldSet, ParticleSet
+
 PACKAGE = "Parcels"
 
 
@@ -61,6 +68,77 @@ def patch_docstring(obj: Callable, extra: str) -> None:
     obj.__doc__ = f"{obj.__doc__ or ''}{extra}".strip()
 
 
+def field_repr(field: Field) -> str:
+    """Return a pretty repr for Field"""
+    out = f"""<{type(field).__name__}>
+    name            : {field.name!r}
+    grid            : {field.grid!r}
+    extrapolate time: {field.allow_time_extrapolation!r}
+    time_periodic   : {field.time_periodic!r}
+    gridindexingtype: {field.gridindexingtype!r}
+    to_write        : {field.to_write!r}
+"""
+    return textwrap.dedent(out).strip()
+
+
+def _format_list_items_multiline(items: list[str], level: int = 1) -> str:
+    """Given a list of strings, formats them across multiple lines.
+
+    Uses indentation levels of 4 spaces provided by ``level``.
+
+    Example
+    -------
+    >>> output = _format_list_items_multiline(["item1", "item2", "item3"], 4)
+    >>> f"my_items: {output}"
+    my_items: [
+        item1,
+        item2,
+        item3,
+    ]
+    """
+    if len(items) == 0:
+        return "[]"
+
+    assert level >= 1, "Indentation level >=1 supported"
+    indentation_str = level * 4 * " "
+    indentation_str_end = (level - 1) * 4 * " "
+
+    items_str = ",\n".join([textwrap.indent(i, indentation_str) for i in items])
+    return f"[\n{items_str}\n{indentation_str_end}]"
+
+
+def particleset_repr(pset: ParticleSet) -> str:
+    """Return a pretty repr for ParticleSet"""
+    if len(pset) < 10:
+        particles = [repr(p) for p in pset]
+    else:
+        particles = [repr(pset[i]) for i in range(7)] + ["..."]
+
+    out = f"""<{type(pset).__name__}>
+    fieldset   : {pset.fieldset}
+    pclass     : {pset.pclass}
+    repeatdt   : {pset.repeatdt}
+    # particles: {len(pset)}
+    particles  : {_format_list_items_multiline(particles, level=2)}
+"""
+    return textwrap.dedent(out).strip()
+
+
+def fieldset_repr(fieldset: FieldSet) -> str:
+    """Return a pretty repr for FieldSet"""
+    fields_repr = "\n".join([repr(f) for f in fieldset.get_fields()])
+
+    out = f"""<{type(fieldset).__name__}>
+    fields:
+{textwrap.indent(fields_repr, 8 * " ")}
+"""
+    return textwrap.dedent(out).strip()
+
+
+def default_repr(obj: Any):
+    return object.__repr__(obj)
+
+
 def timedelta_to_float(dt: float | timedelta | np.timedelta64) -> float:
     """Convert a timedelta to a float in seconds."""
     if isinstance(dt, timedelta):

parcels/tools/_helpers.py

@@ -1,12 +1,16 @@
+import re
+from datetime import timedelta
 from typing import Any
 
 import numpy as np
 
-from parcels import Grid, TimeConverter
+import parcels
+from parcels import Grid, ParticleFile, TimeConverter, Variable
 from parcels.grid import RectilinearGrid
+from tests.utils import create_fieldset_unit_mesh, create_simple_pset
 
 
-def validate_simple_repr(class_: type, kwargs: dict[str, Any]):
+def assert_simple_repr(class_: type, kwargs: dict[str, Any]):
     """Test that the repr of an object contains all the arguments. This only works for objects where the repr matches the calling signature."""
     obj = class_(**kwargs)
     obj_repr = repr(obj)
@@ -17,12 +21,48 @@ def validate_simple_repr(class_: type, kwargs: dict[str, Any]):
     assert class_.__name__ in obj_repr
 
 
+def valid_indentation(s: str) -> bool:
+    """Make sure that all lines in string is indented with a multiple of 4 spaces."""
+    if s.startswith(" "):
+        return False
+
+    lines = s.split("\n")
+    for line in lines:
+        line = re.sub("^( {4})+", "", line)
+        if line.startswith(" "):
+            return False
+    return True
+
+
+def test_check_indentation():
+    valid = """
+test
+    test
+test
+    test
+        test
+    test"""
+    assert valid_indentation(valid)
+    invalid = """
+test
+    test
+ invalid!
+"""
+    assert not valid_indentation(invalid)
+
+
 def test_grid_repr():
     """Test arguments are in the repr of a Grid object"""
     kwargs = dict(
         lon=np.array([1, 2, 3]), lat=np.array([4, 5, 6]), time=None, time_origin=TimeConverter(), mesh="spherical"
     )
-    validate_simple_repr(Grid, kwargs)
+    assert_simple_repr(Grid, kwargs)
+
+
+def test_variable_repr():
+    """Test arguments are in the repr of the Variable object."""
+    kwargs = dict(name="test", dtype=np.float32, initial=0, to_write=False)
+    assert_simple_repr(Variable, kwargs)
 
 
 def test_rectilineargrid_repr():
@@ -34,4 +74,41 @@ def test_rectilineargrid_repr():
     kwargs = dict(
         lon=np.array([1, 2, 3]), lat=np.array([4, 5, 6]), time=None, time_origin=TimeConverter(), mesh="spherical"
     )
-    validate_simple_repr(RectilinearGrid, kwargs)
+    assert_simple_repr(RectilinearGrid, kwargs)
+
+
+def test_particlefile_repr():
+    pset = create_simple_pset()
+    kwargs = dict(
+        name="file.zarr", particleset=pset, outputdt=timedelta(hours=1), chunks=None, create_new_zarrfile=False
+    )
+    assert_simple_repr(ParticleFile, kwargs)
+
+
+def test_field_repr():
+    field = create_fieldset_unit_mesh().U
+    assert valid_indentation(repr(field))
+
+
+def test_vectorfield_repr():
+    field = create_fieldset_unit_mesh().UV
+    assert isinstance(field, parcels.VectorField)
+    assert valid_indentation(repr(field))
+
+
+def test_fieldset_repr():
+    fieldset = create_fieldset_unit_mesh()
+    assert valid_indentation(repr(fieldset))
+
+
+def test_particleset_repr():
+    pset = create_simple_pset()
+    valid_indentation(repr(pset))
+
+    pset = create_simple_pset(n=15)
+    valid_indentation(repr(pset))
+
+
+def capture(s):
+    with open("file.txt", "a") as f:
+        f.write(s)

tests/test_reprs.py

@@ -3,9 +3,20 @@
 import numpy as np
 import pytest
 
+import parcels.tools._helpers as helpers
 from parcels.tools._helpers import deprecated, deprecated_made_private, timedelta_to_float
 
 
+def test_format_list_items_multiline():
+    expected = """[
+    item1,
+    item2,
+    item3
+]"""
+    assert helpers._format_list_items_multiline(["item1", "item2", "item3"], 1) == expected
+    assert helpers._format_list_items_multiline([], 1) == "[]"
+
+
 def test_deprecated():
     class SomeClass:
         @deprecated()

tests/tools/test_helpers.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 
+import parcels
 from parcels import FieldSet
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
@@ -47,6 +48,18 @@ def create_fieldset_zeros_conversion(mesh="spherical", xdim=200, ydim=100, mesh_
     return FieldSet.from_data(data, dimensions, mesh=mesh)
 
 
+def create_simple_pset(n=1):
+    zeros = np.zeros(n)
+    return parcels.ParticleSet(
+        fieldset=create_fieldset_unit_mesh(),
+        pclass=parcels.ScipyParticle,
+        lon=zeros,
+        lat=zeros,
+        depth=zeros,
+        time=zeros,
+    )
+
+
 def create_spherical_positions(n_particles, max_depth=100000):
     yrange = 2 * np.random.rand(n_particles)
     lat = 180 * (np.arccos(1 - yrange) - 0.5 * np.pi) / np.pi