python-arango-async pyright example

test.py

import json
import pandas as pd
import pydantic
import pydantic_core
from typing import Sequence, cast
from arangoasync.collection import StandardCollection
from arangoasync.database import StandardDatabase
from arangoasync.exceptions import DeserializationError, SerializationError
from arangoasync.serialization import Serializer, Deserializer
from arangoasync.typings import Json, Jsons


class CustomJsonSerializer(Serializer[Json]):
    def dumps(self, data: Json | Sequence[str | Json]) -> str:
        try:
            return json.dumps(data, separators=(",", ":"))
        except Exception as e:
            raise SerializationError("Failed to serialize data to JSON.") from e


class CustomJsonDeserializer(Deserializer[Json, Jsons]):
    def loads(self, data: bytes) -> Json:
        try:
            return json.loads(data)  # type: ignore[no-any-return]
        except Exception as e:
            raise DeserializationError("Failed to deserialize data from JSON.") from e

    def loads_many(self, data: bytes) -> Jsons:
        return self.loads(data)  # type: ignore[return-value]


class Student(pydantic.BaseModel):
    name: str
    age: int


class StudentSerializer(Serializer[Student]):
    def dumps(self, data: Student | Sequence[Student | str]) -> str:
        try:
            if isinstance(data, Student):
                return data.model_dump_json()
            else:
                serialized_data = []
                for student in data:
                    if isinstance(student, str):
                        serialized_data.append(student)
                    else:
                        serialized_data.append(student.model_dump())
                return json.dumps(serialized_data, separators=(",", ":"))
        except Exception as e:
            raise SerializationError("Failed to serialize data.") from e


class StudentDeserializer(Deserializer[Student, pd.DataFrame]):
    def loads(self, data: bytes) -> Student:
        try:
            return Student.model_validate(pydantic_core.from_json(data))
        except Exception as e:
            raise DeserializationError("Failed to deserialize data.") from e

    def loads_many(self, data: bytes) -> pd.DataFrame:
        return pd.DataFrame(json.loads(data))


async def main():
    from arangoasync import ArangoClient
    from arangoasync.auth import Auth

    # Initialize the client for ArangoDB.
    async with ArangoClient(
            hosts="http://localhost:8529",
            serializer=CustomJsonSerializer(),
            deserializer=CustomJsonDeserializer(),
    ) as client:
        auth = Auth(username="root", password="passwd")

        # Connect to "test" database as root user.
        db: StandardDatabase = await client.db("_system", auth=auth, verify=True)

        # Create the "students" collection.
        if not await db.has_collection("students"):
            await db.create_collection("students")
        col = cast(
            StandardCollection[Student, Student, pd.DataFrame],
            db.collection(
                "students",
                doc_serializer=StudentSerializer(),
                doc_deserializer=StudentDeserializer()),
        )

        doc = cast(Json, await col.insert(Student(name="John Doe", age=20)))
        docs = cast(Jsons, await col.insert_many([
            Student(name="Jane Doe", age=22),
            Student(name="Alice Smith", age=19),
            Student(name="Bob Johnson", age=21),
        ]))

        # Get one document.
        john = await col.get(doc)
        assert type(john) == Student

        # Get multiple documents.
        keys = [doc["_key"] for doc in docs]
        students = await col.get_many(keys)
        assert type(students) == pd.DataFrame


if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

Validate with pyright test.py