types: add trajan.Dataset typed alias with .traj accessor

Ship trajan/__init__.pyi so LSPs (basedpyright, pyright, pylance, mypy)
discover Dataset as a typed subclass of xr.Dataset with traj: Traj.

Users annotate once and get full completion on all resulting datasets:

    ds: trajan.Dataset = xr.open_dataset('file.nc')
    ds.traj.speed()  # <- LSP completion works

At runtime Dataset = xr.Dataset (no overhead). The .pyi stub is
authoritative for type checkers and takes precedence over the source.

Also add 'from __future__ import annotations' to traj.py and guard
'from . import Dataset' behind TYPE_CHECKING to break the circular
import that arises when traj.py uses Dataset in return annotations.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: gauteh

trajan/__init__.py

@@ -17,8 +17,15 @@
 
 from . import waves as _
 
+from .traj import Traj
+
+# Runtime alias so `trajan.Dataset` is usable as a value (e.g. isinstance checks).
+# The `.pyi` stub declares the typed subclass with `.traj` typed as Traj.
+Dataset = xr.Dataset
+
 logger = logging.getLogger(__name__)
 
+
 __version__ = importlib.metadata.version("trajan")
 
 def versions():

trajan/__init__.pyi

@@ -0,0 +1,44 @@
+"""
+Stub file for trajan public API.
+
+Declares ``trajan.Dataset`` — a typed alias for ``xr.Dataset`` that exposes
+the ``.traj`` accessor.  Use it as a type annotation to get LSP completion:
+
+    import trajan
+    ds: trajan.Dataset = xr.open_dataset("file.nc")
+    ds.traj.speed()   # <- full completion
+"""
+
+import pandas as pd
+import xarray as xr
+from typing import Any
+
+from .traj import Traj as Traj
+from .traj1d import Traj1d as Traj1d
+from .traj2d import Traj2d as Traj2d
+
+class Dataset(xr.Dataset):
+    """xarray Dataset with the trajan ``.traj`` accessor typed."""
+
+    @property
+    def traj(self) -> Traj: ...
+
+def versions() -> str: ...
+
+def read_csv(f: Any, **kwargs: Any) -> Dataset: ...
+
+def from_dataframe(
+    df: pd.DataFrame,
+    lon: str = ...,
+    lat: str = ...,
+    time: str = ...,
+    name: str | None = ...,
+    *,
+    __test_condense__: bool = ...,
+) -> Dataset: ...
+
+def trajectory_dict_to_dataset(
+    trajectory_dict: dict[str, Any],
+    variable_attributes: dict[str, Any] | None = ...,
+    global_attributes: dict[str, Any] | None = ...,
+) -> Dataset: ...

trajan/traj.py

@@ -5,10 +5,13 @@
 https://cfconventions.org/Data/cf-conventions/cf-conventions-1.10/cf-conventions.html#_multidimensional_array_representation_of_trajectories.
 """
 
+from __future__ import annotations
+
 from abc import abstractmethod
 from datetime import timedelta
 from functools import cache
 import inspect
+from typing import TYPE_CHECKING
 import pyproj
 import numpy as np
 import xarray as xr
@@ -20,6 +23,9 @@
 from .plot import Plot
 from .animation import Animation
 
+if TYPE_CHECKING:
+    from . import Dataset
+
 logger = logging.getLogger(__name__)
 
 
@@ -101,7 +107,7 @@ def grid_area(lons, lats):
 
 
 class Traj:
-    ds: xr.Dataset
+    ds: Dataset
 
     __plot__: Plot
     __animate__: Animation
@@ -468,7 +474,7 @@ def ccrs(self) -> cartopy.crs.CRS:
         else:
             return None
 
-    def set_crs(self, crs) -> xr.Dataset:
+    def set_crs(self, crs) -> Dataset:
         """
         Returns a new dataset with the CF-supported grid-mapping / projection set to `crs`.
 
@@ -570,7 +576,7 @@ def assign_cf_attrs(self,
                         creator_email=None,
                         title=None,
                         summary=None,
-                        **kwargs) -> xr.Dataset:
+                        **kwargs) -> Dataset:
         """
         Return a new dataset with CF-standard and common attributes set.
 
@@ -724,7 +730,7 @@ def velocity_spectrum(self) -> xr.DataArray:
     #     pass
 
     @abstractmethod
-    def distance_to(self, other) -> xr.Dataset:
+    def distance_to(self, other) -> Dataset:
         """
         Distance between trajectories or a single point.
 
@@ -1002,7 +1008,7 @@ def get_area_convex_hull(self):
                             attrs={"units": "m2"})
 
     @abstractmethod
-    def gridtime(self, times, time_varname=None) -> xr.Dataset:
+    def gridtime(self, times, time_varname=None) -> Dataset:
         """Interpolate dataset to a regular time interval or a different grid.
 
         Parameters
@@ -1022,7 +1028,7 @@ def gridtime(self, times, time_varname=None) -> xr.Dataset:
         """
 
     @abstractmethod
-    def sel(self, *args, **kwargs) -> xr.Dataset:
+    def sel(self, *args, **kwargs) -> Dataset:
         """Select on each trajectory. On 1D datasets this is just a shortcut for `Dataset.sel`.
 
         Parameters
@@ -1041,7 +1047,7 @@ def sel(self, *args, **kwargs) -> xr.Dataset:
         """
 
     @abstractmethod
-    def seltime(self, t0=None, t1=None) -> xr.Dataset:
+    def seltime(self, t0=None, t1=None) -> Dataset:
         """Select observations in time window between `t0` and `t1` (inclusive). For 1D datasets prefer to use `xarray.Dataset.sel`.
 
         Parameters
@@ -1060,7 +1066,7 @@ def seltime(self, t0=None, t1=None) -> xr.Dataset:
         """
 
     @abstractmethod
-    def iseltime(self, i) -> xr.Dataset:
+    def iseltime(self, i) -> Dataset:
         """Select observations by index (of non-nan, time, observation) across
         trajectories. For 1D datasets prefer to use `xarray.Dataset.isel`.
 
@@ -1142,7 +1148,7 @@ def iseltime(self, i) -> xr.Dataset:
         """
 
     @abstractmethod
-    def skill(self, expected, method='liu-weissberg', **kwargs) -> xr.Dataset:
+    def skill(self, expected, method='liu-weissberg', **kwargs) -> Dataset:
         """
         Compare the skill score between this trajectory and an `expected` trajectory.
 
@@ -1209,7 +1215,7 @@ def skill(self, expected, method='liu-weissberg', **kwargs) -> xr.Dataset:
         """
 
     @abstractmethod
-    def condense_obs(self) -> xr.Dataset:
+    def condense_obs(self) -> Dataset:
         """
         Move all observations to the first index, so that the observation
         dimension is reduced to a minimum. When creating ragged arrays the
@@ -1234,14 +1240,14 @@ def condense_obs(self) -> xr.Dataset:
         """
 
     @abstractmethod
-    def to_1d(self) -> xr.Dataset:
+    def to_1d(self) -> Dataset:
         """
         Convert dataset into a 1D dataset from. This is only possible if the
         dataset has a single trajectory.
         """
 
     @abstractmethod
-    def to_2d(self, obs_dim='obs') -> xr.Dataset:
+    def to_2d(self, obs_dim='obs') -> Dataset:
         """Convert the dataset to a 2D representation.
 
         Parameters
@@ -1256,13 +1262,13 @@ def to_2d(self, obs_dim='obs') -> xr.Dataset:
         """
 
     @abstractmethod
-    def append(self, da, obs_dims=None) -> xr.Dataset:
+    def append(self, da, obs_dims=None) -> Dataset:
         """
         Append trajectories from other dataset to this.
         """
 
     @abstractmethod
-    def filter(self, method='speed', max_speed=10., nsigma=5.0, side_half_width=2) -> xr.Dataset:
+    def filter(self, method='speed', max_speed=10., nsigma=5.0, side_half_width=2) -> Dataset:
         """Filter outlier positions from trajectories.
 
         Parameters