Adding reprs for FieldSet, Field, VectorField and XGrid

Author: erikvansebille

src/parcels/_core/field.py

@@ -24,7 +24,7 @@
 from parcels._core.uxgrid import UxGrid
 from parcels._core.xgrid import XGrid, _transpose_xfield_data_to_tzyx, assert_all_field_dims_have_axis
 from parcels._python import assert_same_function_signature
-from parcels._reprs import default_repr
+from parcels._reprs import field_repr, vectorfield_repr
 from parcels._typing import VectorType
 from parcels.interpolators import (
     ZeroInterpolator,
@@ -148,6 +148,9 @@ def __init__(
             if "time" not in self.data.coords:
                 raise ValueError("Field data is missing a 'time' coordinate.")
 
+    def __repr__(self):
+        return field_repr(self)
+
     @property
     def units(self):
         return self._units
@@ -272,11 +275,7 @@ def __init__(
             self._vector_interp_method = vector_interp_method
 
     def __repr__(self):
-        return f"""<{type(self).__name__}>
-    name: {self.name!r}
-    U: {default_repr(self.U)}
-    V: {default_repr(self.V)}
-    W: {default_repr(self.W)}"""
+        return vectorfield_repr(self)
 
     @property
     def vector_interp_method(self):

src/parcels/_core/fieldset.py

@@ -18,6 +18,7 @@
 from parcels._core.uxgrid import UxGrid
 from parcels._core.xgrid import _DEFAULT_XGCM_KWARGS, XGrid
 from parcels._logger import logger
+from parcels._reprs import fieldset_repr
 from parcels._typing import Mesh
 from parcels.interpolators import UxPiecewiseConstantFace, UxPiecewiseLinearNode, XConstantField, XLinear
 
@@ -75,6 +76,9 @@ def __getattr__(self, name):
         else:
             raise AttributeError(f"FieldSet has no attribute '{name}'")
 
+    def __repr__(self):
+        return fieldset_repr(self)
+
     @property
     def time_interval(self):
         """Returns the valid executable time interval of the FieldSet,

src/parcels/_core/xgrid.py

@@ -9,6 +9,7 @@
 
 from parcels._core.basegrid import BaseGrid
 from parcels._core.index_search import _search_1d_array, _search_indices_curvilinear_2d
+from parcels._reprs import xgrid_repr
 from parcels._typing import assert_valid_mesh
 
 _XGRID_AXES = Literal["X", "Y", "Z"]
@@ -135,6 +136,9 @@ def from_dataset(cls, ds: xr.Dataset, mesh, xgcm_kwargs=None):
         grid = xgcm.Grid(ds, **xgcm_kwargs)
         return cls(grid, mesh=mesh)
 
+    def __repr__(self):
+        return xgrid_repr(self)
+
     @property
     def axes(self) -> list[_XGRID_AXES]:
         return _get_xgrid_axes(self.xgcm_grid)

src/parcels/_reprs.py

@@ -5,16 +5,67 @@
 import textwrap
 from typing import TYPE_CHECKING, Any
 
+import xarray as xr
+
 if TYPE_CHECKING:
     from parcels import Field, FieldSet, ParticleSet
 
 
-def field_repr(field: Field) -> str:  # TODO v4: Rework or remove entirely
+def fieldset_repr(fieldset: FieldSet) -> str:
+    """Return a pretty repr for FieldSet"""
+    fields = [f for f in fieldset.fields.values() if getattr(f.__class__, "__name__", "") == "Field"]
+    vfields = [f for f in fieldset.fields.values() if getattr(f.__class__, "__name__", "") == "VectorField"]
+
+    fields_repr = "\n".join([repr(f) for f in fields])
+    vfields_repr = "\n".join([vectorfield_repr(vf, from_fieldset_repr=True) for vf in vfields])
+
+    out = f"""<{type(fieldset).__name__}>
+    fields:
+{textwrap.indent(fields_repr, 8 * " ")}
+    vectorfields:
+{textwrap.indent(vfields_repr, 8 * " ")}
+"""
+    return textwrap.dedent(out).strip()
+
+
+def field_repr(field: Field, offset: int = 0) -> str:
     """Return a pretty repr for Field"""
-    out = f"""<{type(field).__name__}>
-    name            : {field.name!r}
-    data            : {field.data!r}
-    extrapolate time: {field.allow_time_extrapolation!r}
+    with xr.set_options(display_expand_data=False):
+        out = f"""<{type(field).__name__} {field.name!r}>
+    Parcels attributes:
+        name            : {field.name!r}
+        interp_method   : {field.interp_method!r}
+        time_interval   : {field.time_interval!r}
+        units           : {field.units!r}
+        igrid           : {field.igrid!r}
+    DataArray:
+{textwrap.indent(repr(field.data), 8 * " ")}
+{textwrap.indent(repr(field.grid), 4 * " ")}
+"""
+    return textwrap.indent(out, " " * offset).strip()
+
+
+def vectorfield_repr(fieldset: FieldSet, from_fieldset_repr=False) -> str:
+    """Return a pretty repr for VectorField"""
+    out = f"""<{type(fieldset).__name__} {fieldset.name!r}>
+    Parcels attributes:
+        name                  : {fieldset.name!r}
+        vector_interp_method  : {fieldset.vector_interp_method!r}
+        vector_type           : {fieldset.vector_type!r}
+    {field_repr(fieldset.U, offset=4) if not from_fieldset_repr else ""}
+    {field_repr(fieldset.V, offset=4) if not from_fieldset_repr else ""}
+    {field_repr(fieldset.W, offset=4) if not from_fieldset_repr and fieldset.W else ""}"""
+    return out
+
+
+def xgrid_repr(grid: Any) -> str:
+    """Return a pretty repr for Grid"""
+    out = f"""<{type(grid).__name__}>
+    Parcels attributes:
+        mesh                  : {grid._mesh}
+        spatialhash           : {grid._spatialhash}
+    xgcm Grid:
+{textwrap.indent(repr(grid.xgcm_grid), 8 * " ")}
 """
     return textwrap.dedent(out).strip()
 
@@ -62,17 +113,6 @@ def particleset_repr(pset: ParticleSet) -> str:
     return textwrap.dedent(out).strip()
 
 
-def fieldset_repr(fieldset: FieldSet) -> str:  # TODO v4: Rework or remove entirely
-    """Return a pretty repr for FieldSet"""
-    fields_repr = "\n".join([repr(f) for f in fieldset.fields.values()])
-
-    out = f"""<{type(fieldset).__name__}>
-    fields:
-{textwrap.indent(fields_repr, 8 * " ")}
-"""
-    return textwrap.dedent(out).strip()
-
-
 def default_repr(obj: Any):
     if is_builtin_object(obj):
         return repr(obj)