Add typing support for get | filters | all methods when using Pydantic

Author: dgarros

docs/docs/python-sdk/examples/schema_pydantic.py renamed to docs/docs/python-sdk/examples/pydantic_car.py

@@ -4,36 +4,47 @@
 
 from typing import Annotated
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ConfigDict
 from infrahub_sdk import InfrahubClient
 from rich import print as rprint
-from infrahub_sdk.schema import InfrahubAttributeParam as AttrParam,  InfrahubRelationshipParam as RelParam, AttributeKind, from_pydantic
+from infrahub_sdk.schema import InfrahubAttributeParam as AttrParam,  InfrahubRelationshipParam as RelParam, AttributeKind, from_pydantic, NodeSchema, NodeModel, GenericModel
 
 
-class Tag(BaseModel):
+class Tag(NodeModel):
+    model_config = ConfigDict(
+        node_schema=NodeSchema(name="Tag", namespace="Test", human_readable_fields=["name__value"])
+    )
+    
     name: Annotated[str, AttrParam(unique=True), Field(description="The name of the tag")]
     label: str | None = Field(description="The label of the tag")
     description: Annotated[str | None, AttrParam(kind=AttributeKind.TEXTAREA)] = None
 
 
-class Car(BaseModel):
+class TestCar(NodeModel):
     name: str = Field(description="The name of the car")
     tags: list[Tag]
-    owner: Annotated[Person, RelParam(identifier="car__person")]
-    secondary_owner: Person | None = None
+    owner: Annotated[TestPerson, RelParam(identifier="car__person")]
+    secondary_owner: TestPerson | None = None
 
 
-class Person(BaseModel):
+class TestPerson(GenericModel):
     name: str
-    cars: Annotated[list[Car] | None, RelParam(identifier="car__person")] = None
+
+class TestCarOwner(NodeModel, TestPerson):
+    cars: Annotated[list[TestCar] | None, RelParam(identifier="car__person")] = None
 
 
 async def main():
     client = InfrahubClient()
-    schema = from_pydantic(models=[Person, Car, Tag])
+    schema = from_pydantic(models=[TestPerson, TestCar, Tag, TestPerson, TestCarOwner])
     rprint(schema.to_schema_dict())
     response = await client.schema.load(schemas=[schema.to_schema_dict()], wait_until_converged=True)
     rprint(response)
 
+    # Create a Tag
+    tag = await client.create("TestTag", name="Blue", label="Blue")
+    await tag.save(allow_upsert=True)
+
+
 if __name__ == "__main__":
     aiorun(main())

docs/docs/python-sdk/examples/pydantic_infra.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from asyncio import run as aiorun
+
+from infrahub_sdk.async_typer import AsyncTyper
+
+from typing import Annotated
+
+from pydantic import BaseModel, Field, ConfigDict
+from infrahub_sdk import InfrahubClient
+from rich import print as rprint
+from infrahub_sdk.schema import InfrahubAttributeParam as AttrParam,  InfrahubRelationshipParam as RelParam, AttributeKind, from_pydantic, NodeSchema, NodeModel, GenericSchema, GenericModel, RelationshipKind
+
+
+app = AsyncTyper()
+
+
+class Site(NodeModel):
+    model_config = ConfigDict(
+        node_schema=NodeSchema(name="Site", namespace="Infra", human_friendly_id=["name__value"],  display_labels=["name__value"])
+    )
+
+    name: Annotated[str, AttrParam(unique=True)] = Field(description="The name of the site")
+
+
+class Vlan(NodeModel):
+    model_config = ConfigDict(
+        node_schema=NodeSchema(name="Vlan", namespace="Infra", human_friendly_id=["vlan_id__value"],  display_labels=["vlan_id__value"])
+    )
+
+    name: str
+    vlan_id: int
+    description: str | None = None
+
+
+class Device(NodeModel):
+    model_config = ConfigDict(
+        node_schema=NodeSchema(name="Device", namespace="Infra", human_friendly_id=["name__value"], display_labels=["name__value"])
+    )
+
+    name: Annotated[str, AttrParam(unique=True)] = Field(description="The name of the car")
+    site: Annotated[Site, RelParam(kind=RelationshipKind.ATTRIBUTE, identifier="device__site")]
+    interfaces: Annotated[list[Interface], RelParam(kind=RelationshipKind.COMPONENT, identifier="device__interfaces")] = Field(default_factory=list)
+
+
+class Interface(GenericModel):
+    model_config = ConfigDict(
+        generic_schema=GenericSchema(name="Interface", namespace="Infra", human_friendly_id=["device__name__value", "name__value"], display_labels=["name__value"])
+    )
+
+    device: Annotated[Device, RelParam(kind=RelationshipKind.PARENT, identifier="device__interfaces")]
+    name: str
+    description: str | None = None
+
+class L2Interface(Interface):
+    model_config = ConfigDict(
+        node_schema=NodeSchema(name="L2Interface", namespace="Infra")
+    )
+    
+    vlans: list[Vlan] = Field(default_factory=list)
+
+class LoopbackInterface(Interface):
+    model_config = ConfigDict(
+        node_schema=NodeSchema(name="LoopbackInterface", namespace="Infra")
+    )
+    
+
+
+@app.command()
+async def load_schema():
+    client = InfrahubClient()
+    schema = from_pydantic(models=[Site, Device, Interface, L2Interface, LoopbackInterface, Vlan])
+    rprint(schema.to_schema_dict())
+    response = await client.schema.load(schemas=[schema.to_schema_dict()], wait_until_converged=True)
+    rprint(response)
+
+
+@app.command()
+async def load_data():
+    client = InfrahubClient()
+
+    atl = await client.create("InfraSite", name="ATL")
+    await atl.save(allow_upsert=True)
+    cdg = await client.create("InfraSite", name="CDG")
+    await cdg.save(allow_upsert=True)
+
+    device1 = await client.create("InfraDevice", name="atl1-dev1", site=atl)
+    await device1.save(allow_upsert=True)
+    device2 = await client.create("InfraDevice", name="atl1-dev2", site=atl)
+    await device2.save(allow_upsert=True)
+
+    lo0dev1 = await client.create("InfraLoopbackInterface", name="lo0", device=device1)
+    await lo0dev1.save(allow_upsert=True)
+    lo0dev2 = await client.create("InfraLoopbackInterface", name="lo0", device=device2)
+    await lo0dev2.save(allow_upsert=True)
+
+    for idx in range(1, 3):
+        interface = await client.create("InfraL2Interface", name=f"Ethernet{idx}", device=device1)
+        await interface.save(allow_upsert=True)
+
+
+@app.command()
+async def query_data():
+    client = InfrahubClient()
+    sites = await client.all(kind=Site)
+
+    breakpoint()
+    devices = await client.all(kind=Device)
+    for device in devices:
+        rprint(device)
+
+if __name__ == "__main__":
+    app()