From 9d07fc51684651e7f7805c9ffdd8814efcc45ff1 Mon Sep 17 00:00:00 2001 From: Mourits de Beer <31511766+ff137@users.noreply.github.com> Date: Wed, 20 Sep 2023 14:35:22 +0200 Subject: [PATCH] Support both Pydantic v1 and v2 (#65) * unpin to allow pydantic v2 * :art: set default value * :arrow_up: upgrade fastapi-websocket-rpc * version bump * :heavy_plus_sign: add `packaging` as dependency * :art: add util methods for supporting both pydantic 1 and 2 * use `get_model_serializer` to support pydantic v1 and 2 * :art: black formatting * :art: imports * :art: add helper method for printing model dict in pydantic v1 and 2 * :art: support kwargs in helper methods * :art: helper method to print model dict * revert version bump * :art: update pydantic helper methods to return result in one call * :art: update to use new pydantic helper methods --- fastapi_websocket_pubsub/event_broadcaster.py | 19 ++++++++++------- fastapi_websocket_pubsub/event_notifier.py | 3 ++- fastapi_websocket_pubsub/rpc_event_methods.py | 7 ++++++- fastapi_websocket_pubsub/util.py | 21 +++++++++++++++++++ requirements.txt | 7 ++++--- 5 files changed, 45 insertions(+), 12 deletions(-) create mode 100644 fastapi_websocket_pubsub/util.py diff --git a/fastapi_websocket_pubsub/event_broadcaster.py b/fastapi_websocket_pubsub/event_broadcaster.py index 89dadd6..f3e43fe 100644 --- a/fastapi_websocket_pubsub/event_broadcaster.py +++ b/fastapi_websocket_pubsub/event_broadcaster.py @@ -1,12 +1,13 @@ import asyncio -from typing import Any, Union -from pydantic.main import BaseModel -from .event_notifier import EventNotifier, Subscription, TopicList, ALL_TOPICS -from broadcaster import Broadcast +from typing import Any -from .logger import get_logger +from broadcaster import Broadcast from fastapi_websocket_rpc.utils import gen_uid +from pydantic.main import BaseModel +from .event_notifier import ALL_TOPICS, EventNotifier, Subscription, TopicList +from .logger import get_logger +from .util import pydantic_serialize logger = get_logger("EventBroadcaster") @@ -18,7 +19,7 @@ class BroadcastNotification(BaseModel): notifier_id: NotifierId topics: TopicList - data: Any + data: Any = None class EventBroadcasterException(Exception): @@ -180,7 +181,9 @@ async def __broadcast_notifications__(self, subscription: Subscription, data): async with self._broadcast_type( self._broadcast_url ) as sharing_broadcast_channel: - await sharing_broadcast_channel.publish(self._channel, note.json()) + await sharing_broadcast_channel.publish( + self._channel, pydantic_serialize(note) + ) async def _subscribe_to_all_topics(self): return await self._notifier.subscribe( @@ -277,8 +280,10 @@ async def __read_notifications__(self): ) self._tasks.add(task) + def cleanup(task): self._tasks.remove(task) + task.add_done_callback(cleanup) except: logger.exception("Failed handling incoming broadcast") diff --git a/fastapi_websocket_pubsub/event_notifier.py b/fastapi_websocket_pubsub/event_notifier.py index ca0bd61..984cc55 100644 --- a/fastapi_websocket_pubsub/event_notifier.py +++ b/fastapi_websocket_pubsub/event_notifier.py @@ -7,6 +7,7 @@ from pydantic import BaseModel # pylint: disable=no-name-in-module from .logger import get_logger +from .util import pydantic_to_dict logger = get_logger("EventNotifier") @@ -129,7 +130,7 @@ async def subscribe( ) subscriptions.append(new_subscription) new_subscriptions.append(new_subscription) - logger.debug(f"New subscription {new_subscription.dict()}") + logger.debug(f"New subscription {pydantic_to_dict(new_subscription)}") await EventNotifier.trigger_events( self._on_subscribe_events, subscriber_id, topics ) diff --git a/fastapi_websocket_pubsub/rpc_event_methods.py b/fastapi_websocket_pubsub/rpc_event_methods.py index 8265ec0..0238a45 100644 --- a/fastapi_websocket_pubsub/rpc_event_methods.py +++ b/fastapi_websocket_pubsub/rpc_event_methods.py @@ -1,7 +1,10 @@ import asyncio + from fastapi_websocket_rpc import RpcMethodsBase + from .event_notifier import EventNotifier, Subscription, TopicList from .logger import get_logger +from .util import pydantic_to_dict class RpcEventServerMethods(RpcMethodsBase): @@ -22,7 +25,9 @@ async def subscribe(self, topics: TopicList = []) -> bool: async def callback(subscription: Subscription, data): # remove the actual function sub = subscription.copy(exclude={"callback"}) - self.logger.info(f"Notifying other side: subscription={subscription.dict(exclude={'callback'})}, data={data}, channel_id={self.channel.id}") + self.logger.info( + f"Notifying other side: subscription={pydantic_to_dict(subscription, exclude={'callback'})}, data={data}, channel_id={self.channel.id}" + ) await self.channel.other.notify(subscription=sub, data=data) if self._rpc_channel_get_remote_id: diff --git a/fastapi_websocket_pubsub/util.py b/fastapi_websocket_pubsub/util.py new file mode 100644 index 0000000..5e9bd51 --- /dev/null +++ b/fastapi_websocket_pubsub/util.py @@ -0,0 +1,21 @@ +import pydantic +from packaging import version + + +# Helper methods for supporting Pydantic v1 and v2 +def is_pydantic_pre_v2(): + return version.parse(pydantic.VERSION) < version.parse("2.0.0") + + +def pydantic_serialize(model, **kwargs): + if is_pydantic_pre_v2(): + return model.json(**kwargs) + else: + return model.model_dump_json(**kwargs) + + +def pydantic_to_dict(model, **kwargs): + if is_pydantic_pre_v2(): + return model.dict(**kwargs) + else: + return model.model_dump(**kwargs) diff --git a/requirements.txt b/requirements.txt index abfc418..f51121c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ -fastapi-websocket-rpc>=0.1.24,<1 +fastapi-websocket-rpc>=0.1.25,<1 +packaging>=20.4 permit-broadcaster[redis,postgres,kafka]>=0.2.5,<3 -pydantic>=1.9.1,<2 -websockets>=10.3,<11 \ No newline at end of file +pydantic>=1.9.1 +websockets>=10.3,<11