From 7cfb79c7950cc6dacb5a610b33632db3e68d1d76 Mon Sep 17 00:00:00 2001
From: Rosina Derks <rosina_derks@hotmail.com>
Date: Mon, 9 Oct 2023 15:28:36 +0200
Subject: [PATCH] KDP-1528 refactor instances, collections, and dataqueries

---
 src/edr_pydantic/collections.py  |  6 +++---
 src/edr_pydantic/data_queries.py | 19 +++++--------------
 src/edr_pydantic/instances.py    | 31 -------------------------------
 src/edr_pydantic/link.py         |  6 +++++-
 tests/test_edr.py                |  3 +--
 5 files changed, 14 insertions(+), 51 deletions(-)
 delete mode 100644 src/edr_pydantic/instances.py

diff --git a/src/edr_pydantic/collections.py b/src/edr_pydantic/collections.py
index 4d1780e..f32a347 100644
--- a/src/edr_pydantic/collections.py
+++ b/src/edr_pydantic/collections.py
@@ -2,6 +2,7 @@
 from typing import List
 from typing import Optional
 
+from .data_queries import CollectionDataQueries
 from .data_queries import DataQueries
 from .extent import Extent
 from .link import Link
@@ -9,7 +10,6 @@
 from .parameter import Parameter
 
 
-# TODO instances or collections as base?
 class Collection(EDRBaseModel):
     id: str
     title: Optional[str] = None
@@ -17,7 +17,7 @@ class Collection(EDRBaseModel):
     keywords: Optional[List[str]] = None
     links: List[Link]
     extent: Extent
-    data_queries: Optional[DataQueries] = None
+    data_queries: Optional[CollectionDataQueries] = None
     # TODO According to req A.13 it should be CRSDetails object
     crs: Optional[List[str]] = None
     output_formats: Optional[List[str]] = None
@@ -33,7 +33,7 @@ class CollectionsModel(EDRBaseModel):
 
 # For now, the instance metadata corresponds to the first collection metadata. So they have equal classes
 class Instance(Collection):
-    pass
+    data_queries: Optional[DataQueries] = None
 
 
 class InstancesModel(EDRBaseModel):
diff --git a/src/edr_pydantic/data_queries.py b/src/edr_pydantic/data_queries.py
index 9d552d8..8d8791f 100644
--- a/src/edr_pydantic/data_queries.py
+++ b/src/edr_pydantic/data_queries.py
@@ -1,21 +1,9 @@
 from typing import Optional
-from typing import Union
 
+from .link import EDRQueryLink
 from .my_base_model import EDRBaseModel
-from .variables import CorridorVariables
-from .variables import CubeVariables
-from .variables import ItemVariables
-from .variables import RadiusVariables
-from .variables import Variables
 
 
-class EDRQueryLink(EDRBaseModel, extra="allow"):
-    href: str
-    rel: str
-    variables: Union[Variables, CubeVariables, CorridorVariables, ItemVariables, RadiusVariables]
-
-
-# TODO Why not use the Link object, difference is required variables for this link?
 class EDRQuery(EDRBaseModel):
     link: EDRQueryLink
 
@@ -27,7 +15,10 @@ class DataQueries(EDRBaseModel):
     cube: Optional[EDRQuery] = None
     trajectory: Optional[EDRQuery] = None
     corridor: Optional[EDRQuery] = None
-    # TODO Difference object/req: item, location, plus instances (only for collection)?
+    # TODO Difference object/req: item, location?
     locations: Optional[EDRQuery] = None
     items: Optional[EDRQuery] = None
+
+
+class CollectionDataQueries(DataQueries):
     instances: Optional[EDRQuery] = None
diff --git a/src/edr_pydantic/instances.py b/src/edr_pydantic/instances.py
deleted file mode 100644
index 93180cd..0000000
--- a/src/edr_pydantic/instances.py
+++ /dev/null
@@ -1,31 +0,0 @@
-from typing import Dict
-from typing import List
-from typing import Optional
-
-from .data_queries import DataQueries
-from .extent import Extent
-from .link import Link
-from .my_base_model import EDRBaseModel
-from .parameter import Parameter
-
-
-# TODO Difference between Instance and Collection?
-class Instance(EDRBaseModel):
-    id: str
-    title: Optional[str] = None
-    description: Optional[str] = None
-    keywords: Optional[List[str]] = None
-    links: List[Link]
-    extent: Extent
-    data_queries: Optional[DataQueries] = None
-    # TODO According to req A.13 it should be CRSDetails object
-    crs: Optional[List[str]] = None
-    output_formats: Optional[List[str]] = None
-    parameter_names: Dict[str, Parameter]
-    # TODO According to req A.13 MAY have distanceunits. If radius is in link, it SHALL have distanceunits
-    distanceunits: Optional[List[str]] = None
-
-
-class InstancesModel(EDRBaseModel):
-    links: List[Link]
-    instances: List[Instance]
diff --git a/src/edr_pydantic/link.py b/src/edr_pydantic/link.py
index b1020d0..a35c4fa 100644
--- a/src/edr_pydantic/link.py
+++ b/src/edr_pydantic/link.py
@@ -18,5 +18,9 @@ class Link(EDRBaseModel):
     title: Optional[str] = None
     length: Optional[int] = None
     templated: Optional[bool] = None
-    # TODO Dataquery link separate from Link because variables is required there?
     variables: Optional[Union[Variables, CubeVariables, CorridorVariables, ItemVariables, RadiusVariables]] = None
+
+
+# For EDRQueryLink the variables element is required compared to the Link object
+class EDRQueryLink(Link, extra="allow"):
+    variables: Union[Variables, CubeVariables, CorridorVariables, ItemVariables, RadiusVariables]
diff --git a/tests/test_edr.py b/tests/test_edr.py
index 73efab5..c924f2d 100644
--- a/tests/test_edr.py
+++ b/tests/test_edr.py
@@ -2,10 +2,9 @@
 from pathlib import Path
 
 import pytest
+from edr_pydantic.collections import CollectionsModel
 from pydantic import ValidationError
 
-from src.edr_pydantic.collections import CollectionsModel
-
 
 happy_cases = [
     ("knmi-example-collection.json", CollectionsModel),