unifiying "template_table" with just the "template" argument.

Author: alxmrs

tests/test_ds.py

@@ -7,7 +7,7 @@
 * Round-trip identity across varied source Datasets (one parametrized
   ``assert_identical`` test, not eight per-aspect checks).
 * Aggregation, ``dimension_columns`` inference, and the template /
-  ``template_table`` resolution rules with their error paths.
+  ``template`` resolution rules (name or Dataset) with their error paths.
 * Sparsity handling and ``fill_value`` dtype behavior.
 * The vectorized-indexer fallback through xarray's adapter.
 
@@ -267,12 +267,12 @@ def test_chunks_auto_snaps_to_source_partitions():
 def test_to_dataset_multi_registered_requires_explicit_template(
     air_dataset_small,
 ):
-    """With more than one registered Dataset, the caller must
-    disambiguate via ``template_table=``."""
+    """With more than one registered Dataset, the caller disambiguates by
+    passing a registered table name as ``template=``."""
     ctx = XarrayContext()
     ctx.from_dataset("air1", air_dataset_small)
     ctx.from_dataset("air2", air_dataset_small)
-    out = ctx.sql("SELECT * FROM air1").to_dataset(template_table="air1")
+    out = ctx.sql("SELECT * FROM air1").to_dataset(template="air1")
     assert set(out.dims) == {"time", "lat", "lon"}
 
 
@@ -288,36 +288,31 @@ def test_to_dataset_infer_fails_when_no_template_fits(air_dataset_small):
         ).to_dataset()
 
 
-def test_template_table_explicit_override(air_dataset_small):
-    """``template_table=`` picks a registered Dataset deterministically."""
+def test_template_accepts_name_or_dataset(air_dataset_small):
+    """``template=`` accepts either a registered table name or a Dataset
+    object, with equivalent metadata recovery."""
     other = air_dataset_small.copy()
     other.attrs = {"flag": "other"}
     ctx = XarrayContext()
     ctx.from_dataset("air", air_dataset_small)
     ctx.from_dataset("other", other)
-    out = ctx.sql("SELECT * FROM air").to_dataset(
-        dims=["time", "lat", "lon"], template_table="other"
+
+    by_name = ctx.sql("SELECT * FROM air").to_dataset(
+        dims=["time", "lat", "lon"], template="other"
+    )
+    by_object = ctx.sql("SELECT * FROM air").to_dataset(
+        dims=["time", "lat", "lon"], template=other
     )
-    assert out.attrs == {"flag": "other"}
+    assert by_name.attrs == {"flag": "other"}
+    assert by_object.attrs == {"flag": "other"}
 
 
-def test_template_table_unknown_raises(air_dataset_small):
+def test_template_unknown_name_raises(air_dataset_small):
     ctx = XarrayContext()
     ctx.from_dataset("air", air_dataset_small)
     with pytest.raises(ValueError, match="not a registered table"):
         ctx.sql("SELECT * FROM air").to_dataset(
-            dims=["time", "lat", "lon"], template_table="missing"
-        )
-
-
-def test_template_and_template_table_mutually_exclusive(air_dataset_small):
-    ctx = XarrayContext()
-    ctx.from_dataset("air", air_dataset_small)
-    with pytest.raises(ValueError, match="Pass at most one"):
-        ctx.sql("SELECT * FROM air").to_dataset(
-            dims=["time", "lat", "lon"],
-            template=air_dataset_small,
-            template_table="air",
+            dims=["time", "lat", "lon"], template="missing"
         )
 
 

xarray_sql/ds.py

@@ -663,8 +663,7 @@ def to_pandas(self) -> pd.DataFrame:
     def to_dataset(
         self,
         dims: list[str] | None = None,
-        template: xr.Dataset | None = None,
-        template_table: str | None = None,
+        template: xr.Dataset | str | None = None,
         sparsity: Sparsity = "result",
         fill_value: Any = np.nan,
         chunks: Mapping[str, int] | str | None = "inherit",
@@ -677,10 +676,11 @@ def to_dataset(
                 referenced by the SQL ``FROM`` clause (if exactly one
                 matches), or any single registered Dataset whose dims are
                 all present in the result columns.
-            template: Source ``xr.Dataset`` to recover metadata from.
-                Overrides any auto-resolved template.
-            template_table: Name of a registered table to use as the
-                template. Mutually exclusive with ``template``.
+            template: Source to recover metadata (attrs, encoding, non-dim
+                coordinates, dim-coord dtype) from. Either an ``xr.Dataset``
+                used directly, or the name of a registered table (e.g.
+                ``"era5.surface"``) whose Dataset is looked up. When ``None``
+                and exactly one Dataset is registered, that one is used.
             sparsity: ``"result"`` (default) keeps only dim values
                 present in the result. ``"template"`` reindexes to the
                 template's full coord ranges, filling absent cells with
@@ -715,15 +715,13 @@ def to_dataset(
         Raises:
             ValueError: ``dims`` cannot be inferred, names a missing
                 column, or the result has duplicate dim tuples;
-                ``template_table`` is unknown; both ``template`` and
-                ``template_table`` are passed; or
+                ``template`` names an unknown registered table; or
                 ``sparsity="template"`` is requested without a
                 resolvable template.
         """
-        if template is not None and template_table is not None:
-            raise ValueError("Pass at most one of template= or template_table=")
-        if template is None:
-            template = self._resolve_template(template_table)
+        if not isinstance(template, xr.Dataset):
+            # ``template`` is a registered-table name or None; look it up.
+            template = self._resolve_template(template)
         if dims is None:
             dims = self._infer_dimension_columns(preferred_template=template)
         resolved_chunks = self._resolve_chunks(chunks, template, dims)
@@ -770,25 +768,22 @@ def _resolve_chunks(
             return inherited or None
         return chunks
 
-    def _resolve_template(
-        self, template_table: str | None
-    ) -> xr.Dataset | None:
-        """Pick a template Dataset for metadata recovery.
+    def _resolve_template(self, name: str | None) -> xr.Dataset | None:
+        """Pick a template Dataset for metadata recovery by registered name.
 
         Priority:
-          1. Explicit ``template_table`` argument.
+          1. The named registered table (``name``).
           2. If exactly one Dataset is registered on the context, use it.
           3. None.
         """
         templates = self._templates
-        if template_table is not None:
-            if template_table not in templates:
+        if name is not None:
+            if name not in templates:
                 raise ValueError(
-                    f"template_table={template_table!r} is not a "
-                    "registered table on this context. Registered: "
-                    f"{list(templates)}"
+                    f"template={name!r} is not a registered table on this "
+                    f"context. Registered: {list(templates)}"
                 )
-            return templates[template_table]
+            return templates[name]
         if len(templates) == 1:
             return next(iter(templates.values()))
         return None