Fix typehint transformation notification (#1798)

cognitedata · Jun 5, 2024 · 147e9b7 · 147e9b7
1 parent a50067b
commit 147e9b7
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 10 deletions.
diff --git a/cognite/client/_api/transformations/notifications.py b/cognite/client/_api/transformations/notifications.py
@@ -1,14 +1,17 @@
 from __future__ import annotations
 
-from typing import Sequence
+from typing import Sequence, overload
 
 from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ
-from cognite.client.data_classes import TransformationNotification, TransformationNotificationList
+from cognite.client.data_classes import (
+    TransformationNotification,
+    TransformationNotificationList,
+    TransformationNotificationWrite,
+)
 from cognite.client.data_classes.transformations.notifications import (
     TransformationNotificationCore,
     TransformationNotificationFilter,
-    TransformationNotificationWrite,
 )
 from cognite.client.utils._identifier import IdentifierSequence
 from cognite.client.utils._validation import assert_type
@@ -17,13 +20,27 @@
 class TransformationNotificationsAPI(APIClient):
     _RESOURCE_PATH = "/transformations/notifications"
 
+    @overload
+    def create(
+        self, notification: TransformationNotification | TransformationNotificationWrite
+    ) -> TransformationNotification: ...
+
+    @overload
+    def create(
+        self, notification: Sequence[TransformationNotification] | Sequence[TransformationNotificationWrite]
+    ) -> TransformationNotificationList: ...
+
     def create(
-        self, notification: TransformationNotification | Sequence[TransformationNotification]
+        self,
+        notification: TransformationNotification
+        | TransformationNotificationWrite
+        | Sequence[TransformationNotification]
+        | Sequence[TransformationNotificationWrite],
     ) -> TransformationNotification | TransformationNotificationList:
         """`Subscribe for notifications on the transformation errors. <https://developer.cognite.com/api#tag/Transformation-Notifications/operation/createTransformationNotifications>`_
 
         Args:
-            notification (TransformationNotification | Sequence[TransformationNotification]): Notification or list of notifications to create.
+            notification (TransformationNotification | TransformationNotificationWrite | Sequence[TransformationNotification] | Sequence[TransformationNotificationWrite]): Notification or list of notifications to create.
 
         Returns:
             TransformationNotification | TransformationNotificationList: Created notification(s)

diff --git a/cognite/client/data_classes/__init__.py b/cognite/client/data_classes/__init__.py
@@ -269,6 +269,8 @@
 from cognite.client.data_classes.transformations.notifications import (
     TransformationNotification,
     TransformationNotificationList,
+    TransformationNotificationWrite,
+    TransformationNotificationWriteList,
 )
 from cognite.client.data_classes.transformations.schedules import (
     TransformationSchedule,
@@ -480,6 +482,8 @@
     "TransformationJobStatus",
     "TransformationNotification",
     "TransformationNotificationList",
+    "TransformationNotificationWrite",
+    "TransformationNotificationWriteList",
     "TransformationSchedule",
     "TransformationScheduleWrite",
     "TransformationScheduleList",

diff --git a/cognite/client/data_classes/transformations/notifications.py b/cognite/client/data_classes/transformations/notifications.py
@@ -55,11 +55,16 @@ def __init__(
         cognite_client: CogniteClient | None = None,
     ) -> None:
         super().__init__(destination)
-        self.id = id
-        self.transformation_id = transformation_id
-        self.transformation_external_id = transformation_external_id
-        self.created_time = created_time
-        self.last_updated_time = last_updated_time
+        # id/created_time/last_updated_time are required when using the class to read,
+        # but don't make sense passing in when creating a new object. So in order to make the typing
+        # correct here (i.e. int and not Optional[int]), we force the type to be int rather than
+        # Optional[int].
+        # TODO: In the next major version we can make these properties required in the constructor
+        self.id: int = id  # type: ignore
+        self.transformation_id: int = transformation_id  # type: ignore
+        self.transformation_external_id: str = transformation_external_id  # type: ignore
+        self.created_time: int = created_time  # type: ignore
+        self.last_updated_time: int = last_updated_time  # type: ignore
         self._cognite_client = cast("CogniteClient", cognite_client)
 
     def __hash__(self) -> int: