Refactor link-writing into separate module.

tiled/client/base.py

@@ -239,7 +239,6 @@ def new_variation(
             include_data_sources = self._include_data_sources
         return type(self)(
             self.context,
-            item=self._item,
             structure_clients=structure_clients,
             include_data_sources=include_data_sources,
             **kwargs,

tiled/client/union.py

@@ -1,9 +1,56 @@
-from .base import BaseClient
+import copy
+
+from .base import STRUCTURE_TYPES, BaseClient
+from .utils import client_for_item
 
 
 class UnionClient(BaseClient):
     def __repr__(self):
         return (
-            f"<{type(self).__name__} "
-            f"[{', '.join(item.structure_family for item in self.structure().contents)}]>"
+            f"<{type(self).__name__} {{"
+            + ", ".join(f"'{key}'" for key in self.structure().all_keys)
+            + "}>"
+        )
+
+    @property
+    def contents(self):
+        return UnionContents(self)
+
+    def __getitem__(self, key):
+        if key not in self.structure().all_keys:
+            raise KeyError(key)
+        raise NotImplementedError
+
+
+class UnionContents:
+    def __init__(self, node):
+        self.node = node
+
+    def __repr__(self):
+        return (
+            f"<{type(self).__name__} {{"
+            + ", ".join(f"'{item.name}'" for item in self.node.structure().contents)
+            + "}>"
+        )
+
+    def __getitem__(self, name):
+        for index, union_item in enumerate(self.node.structure().contents):
+            if union_item.name == name:
+                structure_family = union_item.structure_family
+                structure_dict = union_item.structure
+                break
+        else:
+            raise KeyError(name)
+        item = copy.deepcopy(self.node.item)
+        item["attributes"]["structure_family"] = structure_family
+        item["attributes"]["structure"] = structure_dict
+        item["links"] = item["links"]["contents"][index]
+        structure_type = STRUCTURE_TYPES[structure_family]
+        structure = structure_type.from_json(structure_dict)
+        return client_for_item(
+            self.node.context,
+            self.node.structure_clients,
+            item,
+            structure=structure,
+            include_data_sources=self.node._include_data_sources,
         )

tiled/server/core.py

@@ -34,6 +34,7 @@
 )
 from . import schemas
 from .etag import tokenize
+from .links import links_for_node
 from .utils import record_timing
 
 del queries
@@ -404,6 +405,7 @@ async def construct_resource(
     depth=0,
 ):
     path_str = "/".join(path_parts)
+    key = path_parts[-1] if path_parts else ""
     attributes = {"ancestors": path_parts[:-1]}
     if include_data_sources and hasattr(entry, "data_sources"):
         attributes["data_sources"] = entry.data_sources
@@ -488,15 +490,17 @@ async def construct_resource(
                         for key, direction in entry.sorting
                     ]
         d = {
-            "id": path_parts[-1] if path_parts else "",
+            "id": key,
             "attributes": schemas.NodeAttributes(**attributes),
         }
         if not omit_links:
-            d["links"] = {
-                "self": f"{base_url}/metadata/{path_str}",
-                "search": f"{base_url}/search/{path_str}",
-                "full": f"{base_url}/container/full/{path_str}",
-            }
+            d["links"] = links_for_node(
+                entry.structure_family,
+                entry.structure(),
+                base_url,
+                "/".join(path_parts[:-1]),
+                key,
+            )
 
         resource = schemas.Resource[
             schemas.NodeAttributes, schemas.ContainerLinks, schemas.ContainerMeta
@@ -510,34 +514,17 @@ async def construct_resource(
                 entry.structure_family
             ]
             links.update(
-                {
-                    link: template.format(base_url=base_url, path=path_str)
-                    for link, template in FULL_LINKS[entry.structure_family].items()
-                }
+                links_for_node(
+                    entry.structure_family,
+                    entry.structure(),
+                    base_url,
+                    "/".join(path_parts[:-1]),
+                    key,
+                )
             )
             structure = asdict(entry.structure())
             if schemas.EntryFields.structure_family in fields:
                 attributes["structure_family"] = entry.structure_family
-            if entry.structure_family == StructureFamily.sparse:
-                shape = structure.get("shape")
-                block_template = ",".join(f"{{{index}}}" for index in range(len(shape)))
-                links[
-                    "block"
-                ] = f"{base_url}/array/block/{path_str}?block={block_template}"
-            elif entry.structure_family == StructureFamily.array:
-                shape = structure.get("shape")
-                block_template = ",".join(
-                    f"{{index_{index}}}" for index in range(len(shape))
-                )
-                links[
-                    "block"
-                ] = f"{base_url}/array/block/{path_str}?block={block_template}"
-            elif entry.structure_family == StructureFamily.table:
-                links[
-                    "partition"
-                ] = f"{base_url}/table/partition/{path_str}?partition={{index}}"
-            elif entry.structure_family == StructureFamily.awkward:
-                links["buffers"] = f"{base_url}/awkward/buffers/{path_str}"
             if schemas.EntryFields.structure in fields:
                 attributes["structure"] = structure
         else:
@@ -719,16 +706,6 @@ class WrongTypeForRoute(Exception):
     pass
 
 
-FULL_LINKS = {
-    StructureFamily.array: {"full": "{base_url}/array/full/{path}"},
-    StructureFamily.awkward: {"full": "{base_url}/awkward/full/{path}"},
-    StructureFamily.container: {"full": "{base_url}/container/full/{path}"},
-    StructureFamily.sparse: {"full": "{base_url}/array/full/{path}"},
-    StructureFamily.table: {"full": "{base_url}/table/full/{path}"},
-    StructureFamily.union: {},
-}
-
-
 def asdict(dc):
     "Compat for converting dataclass or pydantic.BaseModel to dict."
     if dc is None:

tiled/server/links.py

@@ -0,0 +1,90 @@
+"""
+Generate the 'links' section of the response JSON.
+
+The links vary by structure family.
+"""
+from ..structures.core import StructureFamily
+
+
+def links_for_node(structure_family, structure, base_url, path, key):
+    links = {}
+    path_parts = [segment for segment in path.split("/") if segment] + [key]
+    path_str = "/".join(path_parts)
+    links = LINKS_BY_STRUCTURE_FAMILY[structure_family](
+        structure_family, structure, base_url, path_str, key
+    )
+    links["self"] = f"{base_url}/metadata/{path_str}"
+    return links
+
+
+def links_for_array(
+    structure_family, structure, base_url, path_str, key, data_source=None
+):
+    links = {}
+    block_template = ",".join(f"{{{index}}}" for index in range(len(structure.shape)))
+    links["block"] = f"{base_url}/array/block/{path_str}?block={block_template}"
+    links["full"] = f"{base_url}/array/full/{path_str}"
+    if data_source:
+        links["block"] += f"&data_source={data_source}"
+        links["full"] += f"?data_source={data_source}"
+    return links
+
+
+def links_for_awkward(
+    structure_family, structure, base_url, path_str, key, data_source=None
+):
+    links = {}
+    links["buffers"] = f"{base_url}/awkward/buffers/{path_str}"
+    links["full"] = f"{base_url}/awkward/full/{path_str}"
+    if data_source:
+        links["buffers"] += "?data_source={data_source}"
+        links["full"] += "?data_source={data_source}"
+    return links
+
+
+def links_for_container(structure_family, structure, base_url, path_str, key):
+    # Cannot be used inside union, so there is no data_source parameter.
+    links = {}
+    links["full"] = f"{base_url}/container/full/{path_str}"
+    links["search"] = f"{base_url}/search/{path_str}"
+    return links
+
+
+def links_for_table(
+    structure_family, structure, base_url, path_str, key, data_source=None
+):
+    links = {}
+    links["partition"] = f"{base_url}/table/partition/{path_str}?partition={{index}}"
+    links["full"] = f"{base_url}/table/full/{path_str}"
+    if data_source:
+        links["partition"] += f"&data_source={data_source}"
+        links["full"] += f"?data_source={data_source}"
+    return links
+
+
+def links_for_union(structure_family, structure, base_url, path_str, key):
+    links = {}
+    # This contains the links for each structure.
+    links["contents"] = []
+    for item in structure.contents:
+        item_links = LINKS_BY_STRUCTURE_FAMILY[item.structure_family](
+            item.structure_family,
+            item.structure,
+            base_url,
+            path_str,
+            key,
+            data_source=item.name,
+        )
+        item_links["self"] = f"{base_url}/metadata/{path_str}"
+        links["contents"].append(item_links)
+    return links
+
+
+LINKS_BY_STRUCTURE_FAMILY = {
+    StructureFamily.array: links_for_array,
+    StructureFamily.awkward: links_for_awkward,
+    StructureFamily.container: links_for_container,
+    StructureFamily.sparse: links_for_array,  # spare and array are the same
+    StructureFamily.table: links_for_table,
+    StructureFamily.union: links_for_union,
+}

tiled/server/router.py

@@ -45,6 +45,7 @@
     get_validation_registry,
     slice_,
 )
+from .links import links_for_node
 from .settings import get_settings
 from .utils import filter_for_access, get_base_url, record_timing
 
@@ -1175,32 +1176,9 @@ async def _create_node(
         specs=body.specs,
         data_sources=body.data_sources,
     )
-    links = {}
-    base_url = get_base_url(request)
-    path_parts = [segment for segment in path.split("/") if segment] + [key]
-    path_str = "/".join(path_parts)
-    links["self"] = f"{base_url}/metadata/{path_str}"
-    if body.structure_family in {StructureFamily.array, StructureFamily.sparse}:
-        block_template = ",".join(
-            f"{{{index}}}" for index in range(len(node.structure().shape))
-        )
-        links["block"] = f"{base_url}/array/block/{path_str}?block={block_template}"
-        links["full"] = f"{base_url}/array/full/{path_str}"
-    elif body.structure_family == StructureFamily.table:
-        links[
-            "partition"
-        ] = f"{base_url}/table/partition/{path_str}?partition={{index}}"
-        links["full"] = f"{base_url}/table/full/{path_str}"
-    elif body.structure_family == StructureFamily.container:
-        links["full"] = f"{base_url}/container/full/{path_str}"
-        links["search"] = f"{base_url}/search/{path_str}"
-    elif body.structure_family == StructureFamily.awkward:
-        links["buffers"] = f"{base_url}/awkward/buffers/{path_str}"
-        links["full"] = f"{base_url}/awkward/full/{path_str}"
-    elif body.structure_family == StructureFamily.union:
-        pass  # TODO
-    else:
-        raise NotImplementedError(body.structure_family)
+    links = links_for_node(
+        structure_family, structure, get_base_url(request), path, key
+    )
     structure = node.structure()
     if structure is not None:
         structure = structure.dict()

tiled/server/schemas.py

@@ -224,6 +224,9 @@ class SparseLinks(pydantic.BaseModel):
 
 class UnionLinks(pydantic.BaseModel):
     self: str
+    contents: List[
+        Union[ArrayLinks, AwkwardLinks, ContainerLinks, DataFrameLinks, SparseLinks]
+    ]
 
 
 resource_links_type_by_structure_family = {