Fetch more fields automatically

pydrawise/client.py

@@ -81,7 +81,10 @@ async def get_user(self, fetch_zones: bool = True) -> User:
         :param fetch_zones: Not used in this implementation.
         :rtype: User
         """
-        selector = self._schema.Query.me.select(*get_selectors(self._schema, User))
+        skip = [] if fetch_zones else ["controllers.zones"]
+        selector = self._schema.Query.me.select(
+            *get_selectors(self._schema, User, skip)
+        )
         result = await self._query(selector)
         return deserialize(User, result["me"])
 

pydrawise/schema.py

@@ -9,7 +9,7 @@
 from typing import Optional, Union
 
 from apischema.conversions import Conversion
-from apischema.metadata import conversion, skip
+from apischema.metadata import conversion
 
 # The names in this file are from the GraphQL schema and don't always adhere to
 # the naming scheme that pylint expects.
@@ -406,7 +406,7 @@ class Controller:
     )
     online: bool = False
     sensors: list[Sensor] = field(default_factory=list)
-    zones: list[Zone] = field(default_factory=list, metadata=skip(deserialization=True))
+    zones: list[Zone] = field(default_factory=list)
     permitted_program_start_times: list[ProgramStartTime] = field(default_factory=list)
     status: Optional[ControllerStatus] = None
 
@@ -419,9 +419,7 @@ class User:
     customer_id: int = 0
     name: str = ""
     email: str = ""
-    controllers: list[Controller] = field(
-        default_factory=list, metadata=skip(deserialization=True)
-    )
+    controllers: list[Controller] = field(default_factory=list)
 
 
 class Query(ABC):

pydrawise/schema_utils.py

@@ -27,19 +27,22 @@ def deserialize(*args, **kwargs):
 _Field = namedtuple("_Field", ["name", "types"])
 
 
-def _fields(cls) -> Iterator[_Field]:
+def _fields(cls, skip: list[str]) -> Iterator[_Field]:
     """Returns _Field objects for every field on the given dataclass.
 
     :meta private:
     """
     hints = get_type_hints(cls)
     for f in fields(cls):
-        skip_md = f.metadata.get(SKIP_METADATA, None)
-        if skip_md and (skip_md.serialization or skip_md.deserialization):
+        if f.name in skip:
             continue
 
-        conversion_md = f.metadata.get(CONVERSION_METADATA, None)
-        if conversion_md:
+        if (skip_md := f.metadata.get(SKIP_METADATA, None)) and (
+            skip_md.serialization or skip_md.deserialization
+        ):
+            continue
+
+        if conversion_md := f.metadata.get(CONVERSION_METADATA, None):
             yield _Field(f.name, [conversion_md.deserialization.source])
             continue
 
@@ -62,18 +65,24 @@ def _fields(cls) -> Iterator[_Field]:
         yield _Field(f.name, [field_type])
 
 
-def get_selectors(ds: DSLSchema, cls: Type) -> list[DSLField]:
+def get_selectors(
+    ds: DSLSchema, cls: Type, skip_fields: list[str] | None = None
+) -> list[DSLField]:
     """Constructs GraphQL selectors for the given dataclass.
 
     :meta private:
     """
     ret = []
-    for f in _fields(cls):
+    skip_now, skip_later = parse_skip(skip_fields or [])
+    for f in _fields(cls, skip_now):
         dsl_field = getattr(getattr(ds, cls.__name__), f.name)
         if len(f.types) == 1:
             [f_type] = f.types
             if is_dataclass(f_type):
-                ret.append(getattr(dsl_field, "select")(*get_selectors(ds, f_type)))
+                f_skip = skip_later.get(f.name, [])
+                ret.append(
+                    getattr(dsl_field, "select")(*get_selectors(ds, f_type, f_skip))
+                )
             else:
                 ret.append(dsl_field)
         else:
@@ -89,3 +98,21 @@ def get_selectors(ds: DSLSchema, cls: Type) -> list[DSLField]:
                 )
             ret.append(getattr(dsl_field, "select")(*sel_args))
     return ret
+
+
+def parse_skip(skip: list[str] | None = None) -> tuple[list[str], dict[str, list[str]]]:
+    """Converts a flat list of skip fields into (skip_now, skip_later).
+
+    skip_now is a list of fields in the current scope to skip.
+    skip_later is a list of descendant fields to skip.
+
+    :meta private:
+    """
+    now, later = [], {}
+    for item in skip:
+        field, _, descendants = item.partition(".")
+        if descendants:
+            later.setdefault(field, []).append(descendants)
+        else:
+            now.append(field)
+    return now, later

tests/test_client.py

@@ -110,24 +110,60 @@ def zone_json():
     }
 
 
-async def test_get_user(api: Hydrawise, mock_session):
+async def test_get_user(api: Hydrawise, mock_session, controller_json, zone_json):
+    controller_json["zones"] = [zone_json]
     mock_session.execute.return_value = {
         "me": {
             "id": 1234,
             "customerId": 1,
             "name": "My Name",
             "email": "me@asdf.com",
+            "controllers": [controller_json],
         }
     }
     user = await api.get_user()
+    mock_session.execute.assert_awaited_once()
+    [selector] = mock_session.execute.await_args.args
+    query = print_ast(selector)
+    assert "controllers {" in query
+    assert "zones {" in query
+    assert user.id == 1234
+    assert user.name == "My Name"
+    assert user.email == "me@asdf.com"
+    assert len(user.controllers) == 1
+    assert len(user.controllers[0].zones) == 1
+
+
+async def test_get_user_no_zones(api: Hydrawise, mock_session, controller_json):
+    mock_session.execute.return_value = {
+        "me": {
+            "id": 1234,
+            "customerId": 1,
+            "name": "My Name",
+            "email": "me@asdf.com",
+            "controllers": [controller_json],
+        }
+    }
+    user = await api.get_user(fetch_zones=False)
+    mock_session.execute.assert_awaited_once()
+    [selector] = mock_session.execute.await_args.args
+    query = print_ast(selector)
+    assert "controllers {" in query
+    assert "zones {" not in query
     assert user.id == 1234
     assert user.name == "My Name"
     assert user.email == "me@asdf.com"
+    assert len(user.controllers) == 1
+    assert len(user.controllers[0].zones) == 0
 
 
 async def test_get_controllers(api: Hydrawise, mock_session, controller_json):
     mock_session.execute.return_value = {"me": {"controllers": [controller_json]}}
     [controller] = await api.get_controllers()
+    mock_session.execute.assert_awaited_once()
+    [selector] = mock_session.execute.await_args.args
+    query = print_ast(selector)
+    assert "zones {" in query
     assert controller.last_contact_time == datetime(2023, 1, 1, 0, 0, 0)
     assert controller.last_action == datetime(2023, 1, 1, 0, 0, 0)
     assert controller.status.actual_water_time.value == timedelta(minutes=10)
@@ -141,6 +177,7 @@ async def test_get_controller(api: Hydrawise, mock_session, controller_json):
     query = print_ast(selector)
     assert "controller(" in query
     assert "controllerId: 9876" in query
+    assert "zones {" in query
 
     assert controller.last_contact_time == datetime(2023, 1, 1, 0, 0, 0)
     assert controller.last_action == datetime(2023, 1, 1, 0, 0, 0)

tests/test_schema_utils.py

@@ -0,0 +1,7 @@
+from pydrawise import schema_utils
+
+
+def test_parse_skip():
+    skip = ["a", "b.c", "b.d", "e.f.g"]
+    want = ["a"], {"b": ["c", "d"], "e": ["f.g"]}
+    assert schema_utils.parse_skip(skip) == want