MAINT Remove pytz dependency

dev-requirements.txt

@@ -426,7 +426,6 @@ pytimeparse==1.1.8
 pytz==2023.3.post1
     # via
     #   croniter
-    #   flytekit
     #   pandas
 pyyaml==6.0.1
     # via

doc-requirements.txt

@@ -983,7 +983,6 @@ pytimeparse==1.1.8
     # via flytekit
 pytz==2023.3.post1
     # via
-    #   flytekit
     #   great-expectations
     #   mlflow
     #   pandas

flytekit/models/execution.py

@@ -2,14 +2,14 @@
 
 import datetime
 import typing
+from datetime import timezone as _timezone
 from typing import Optional
 
 import flyteidl
 import flyteidl.admin.cluster_assignment_pb2 as _cluster_assignment_pb2
 import flyteidl.admin.execution_pb2 as _execution_pb2
 import flyteidl.admin.node_execution_pb2 as _node_execution_pb2
 import flyteidl.admin.task_execution_pb2 as _task_execution_pb2
-import pytz as _pytz
 
 import flytekit
 from flytekit.models import common as _common_models
@@ -576,12 +576,12 @@ def to_flyte_idl(self):
             outputs=self.outputs.to_flyte_idl() if self.outputs is not None else None,
             abort_metadata=self.abort_metadata.to_flyte_idl() if self.abort_metadata is not None else None,
         )
-        obj.started_at.FromDatetime(self.started_at.astimezone(_pytz.UTC).replace(tzinfo=None))
+        obj.started_at.FromDatetime(self.started_at.astimezone(_timezone.utc).replace(tzinfo=None))
         obj.duration.FromTimedelta(self.duration)
         if self.created_at:
-            obj.created_at.FromDatetime(self.created_at.astimezone(_pytz.UTC).replace(tzinfo=None))
+            obj.created_at.FromDatetime(self.created_at.astimezone(_timezone.utc).replace(tzinfo=None))
         if self.updated_at:
-            obj.updated_at.FromDatetime(self.updated_at.astimezone(_pytz.UTC).replace(tzinfo=None))
+            obj.updated_at.FromDatetime(self.updated_at.astimezone(_timezone.utc).replace(tzinfo=None))
         return obj
 
     @classmethod
@@ -603,13 +603,13 @@ def from_flyte_idl(cls, pb2_object):
             error=error,
             outputs=outputs,
             phase=pb2_object.phase,
-            started_at=pb2_object.started_at.ToDatetime().replace(tzinfo=_pytz.UTC),
+            started_at=pb2_object.started_at.ToDatetime().replace(tzinfo=_timezone.utc),
             duration=pb2_object.duration.ToTimedelta(),
             abort_metadata=abort_metadata,
-            created_at=pb2_object.created_at.ToDatetime().replace(tzinfo=_pytz.UTC)
+            created_at=pb2_object.created_at.ToDatetime().replace(tzinfo=_timezone.utc)
             if pb2_object.HasField("created_at")
             else None,
-            updated_at=pb2_object.updated_at.ToDatetime().replace(tzinfo=_pytz.UTC)
+            updated_at=pb2_object.updated_at.ToDatetime().replace(tzinfo=_timezone.utc)
             if pb2_object.HasField("updated_at")
             else None,
         )

flytekit/models/literals.py

@@ -1,7 +1,7 @@
 from datetime import datetime as _datetime
+from datetime import timezone as _timezone
 from typing import Optional
 
-import pytz as _pytz
 from flyteidl.core import literals_pb2 as _literals_pb2
 from google.protobuf.struct_pb2 import Struct
 
@@ -111,7 +111,7 @@ def datetime(self):
         """
         if self._datetime is None or self._datetime.tzinfo is not None:
             return self._datetime
-        return self._datetime.replace(tzinfo=_pytz.UTC)
+        return self._datetime.replace(tzinfo=_timezone.utc)
 
     @property
     def duration(self):
@@ -149,7 +149,7 @@ def to_flyte_idl(self):
         )
         if self.datetime is not None:
             # Convert to UTC and remove timezone so protobuf behaves.
-            primitive.datetime.FromDatetime(self.datetime.astimezone(_pytz.UTC).replace(tzinfo=None))
+            primitive.datetime.FromDatetime(self.datetime.astimezone(_timezone.utc).replace(tzinfo=None))
         if self.duration is not None:
             primitive.duration.FromTimedelta(self.duration)
         return primitive
@@ -165,7 +165,7 @@ def from_flyte_idl(cls, proto):
             float_value=proto.float_value if proto.HasField("float_value") else None,
             string_value=proto.string_value if proto.HasField("string_value") else None,
             boolean=proto.boolean if proto.HasField("boolean") else None,
-            datetime=proto.datetime.ToDatetime().replace(tzinfo=_pytz.UTC) if proto.HasField("datetime") else None,
+            datetime=proto.datetime.ToDatetime().replace(tzinfo=_timezone.utc) if proto.HasField("datetime") else None,
             duration=proto.duration.ToTimedelta() if proto.HasField("duration") else None,
         )
 

flytekit/models/node_execution.py

@@ -1,8 +1,8 @@
 import datetime
 import typing
+from datetime import timezone as _timezone
 
 import flyteidl.admin.node_execution_pb2 as _node_execution_pb2
-import pytz as _pytz
 
 from flytekit.models import common as _common_models
 from flytekit.models.core import catalog as catalog_models
@@ -195,12 +195,12 @@ def to_flyte_idl(self):
             else None,
             task_node_metadata=self.task_node_metadata.to_flyte_idl() if self.task_node_metadata is not None else None,
         )
-        obj.started_at.FromDatetime(self.started_at.astimezone(_pytz.UTC).replace(tzinfo=None))
+        obj.started_at.FromDatetime(self.started_at.astimezone(_timezone.utc.UTC).replace(tzinfo=None))
         obj.duration.FromTimedelta(self.duration)
         if self.created_at:
-            obj.created_at.FromDatetime(self.created_at.astimezone(_pytz.UTC).replace(tzinfo=None))
+            obj.created_at.FromDatetime(self.created_at.astimezone(_timezone.utc.UTC).replace(tzinfo=None))
         if self.updated_at:
-            obj.updated_at.FromDatetime(self.updated_at.astimezone(_pytz.UTC).replace(tzinfo=None))
+            obj.updated_at.FromDatetime(self.updated_at.astimezone(_timezone.utc.UTC).replace(tzinfo=None))
         return obj
 
     @classmethod
@@ -214,16 +214,20 @@ def from_flyte_idl(cls, p):
             output_uri=p.output_uri if p.HasField("output_uri") else None,
             deck_uri=p.deck_uri,
             error=_core_execution.ExecutionError.from_flyte_idl(p.error) if p.HasField("error") else None,
-            started_at=p.started_at.ToDatetime().replace(tzinfo=_pytz.UTC),
+            started_at=p.started_at.ToDatetime().replace(tzinfo=_timezone.utc.UTC),
             duration=p.duration.ToTimedelta(),
             workflow_node_metadata=WorkflowNodeMetadata.from_flyte_idl(p.workflow_node_metadata)
             if p.HasField("workflow_node_metadata")
             else None,
             task_node_metadata=TaskNodeMetadata.from_flyte_idl(p.task_node_metadata)
             if p.HasField("task_node_metadata")
             else None,
-            created_at=p.created_at.ToDatetime().replace(tzinfo=_pytz.UTC) if p.HasField("created_at") else None,
-            updated_at=p.updated_at.ToDatetime().replace(tzinfo=_pytz.UTC) if p.HasField("updated_at") else None,
+            created_at=p.created_at.ToDatetime().replace(tzinfo=_timezone.utc.UTC)
+            if p.HasField("created_at")
+            else None,
+            updated_at=p.updated_at.ToDatetime().replace(tzinfo=_timezone.utc.UTC)
+            if p.HasField("updated_at")
+            else None,
         )
 
 

plugins/flytekit-airflow/tests/test_agent.py

@@ -1,9 +1,8 @@
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 
 from airflow.operators.python import PythonOperator
 from airflow.sensors.bash import BashSensor
 from airflow.sensors.time_sensor import TimeSensor
-from pytz import UTC
 
 from flytekit import workflow
 
@@ -15,7 +14,9 @@ def py_func():
 
 @workflow
 def wf():
-    sensor = TimeSensor(task_id="fire_immediately", target_time=(datetime.now(tz=UTC) + timedelta(seconds=1)).time())
+    sensor = TimeSensor(
+        task_id="fire_immediately", target_time=(datetime.now(tz=timezone.utc) + timedelta(seconds=1)).time()
+    )
     t3 = BashSensor(task_id="Sensor_succeeds", bash_command="exit 0")
     foo = PythonOperator(task_id="foo", python_callable=py_func)
     sensor >> t3 >> foo

setup.py

@@ -66,7 +66,6 @@
         "python-dateutil>=2.1",
         "python-json-logger>=2.0.0",
         "pytimeparse>=1.1.8,<2.0.0",
-        "pytz",
         "pyyaml!=6.0.0,!=5.4.0,!=5.4.1",  # pyyaml is broken with cython 3: https://github.com/yaml/pyyaml/issues/601
         "requests>=2.18.4,<3.0.0",
         "rich",

tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt

@@ -280,7 +280,6 @@ pytimeparse==1.1.8
     # via flytekit
 pytz==2023.3
     # via
-    #   flytekit
     #   pandas
 pyyaml==6.0.1
     # via

tests/flytekit/unit/models/test_execution.py

@@ -1,7 +1,7 @@
 import datetime
+from datetime import timezone
 
 import pytest
-import pytz
 
 from flytekit.models import common as _common_models
 from flytekit.models import execution as _execution
@@ -19,7 +19,7 @@
 
 
 def test_execution_closure_with_output():
-    test_datetime = datetime.datetime(year=2022, month=1, day=1, tzinfo=pytz.UTC)
+    test_datetime = datetime.datetime(year=2022, month=1, day=1, tzinfo=timezone.utc)
     test_timedelta = datetime.timedelta(seconds=10)
     test_outputs = _execution.LiteralMapBlob(values=_OUTPUT_MAP, uri="http://foo/")
 
@@ -46,7 +46,7 @@ def test_execution_closure_with_output():
 
 
 def test_execution_closure_with_error():
-    test_datetime = datetime.datetime(year=2022, month=1, day=1, tzinfo=pytz.UTC)
+    test_datetime = datetime.datetime(year=2022, month=1, day=1, tzinfo=timezone.utc)
     test_timedelta = datetime.timedelta(seconds=10)
     test_error = _core_exec.ExecutionError(
         code="foo", message="bar", error_uri="http://foobar", kind=_core_exec.ExecutionError.ErrorKind.USER
@@ -75,7 +75,7 @@ def test_execution_closure_with_error():
 
 
 def test_execution_closure_with_abort_metadata():
-    test_datetime = datetime.datetime(year=2022, month=1, day=1, tzinfo=pytz.UTC)
+    test_datetime = datetime.datetime(year=2022, month=1, day=1, tzinfo=timezone.utc)
     test_timedelta = datetime.timedelta(seconds=10)
     abort_metadata = _execution.AbortMetadata(cause="cause", principal="skinner")
 

tests/flytekit/unit/models/test_literals.py

@@ -1,7 +1,6 @@
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 
 import pytest
-import pytz
 
 from flytekit.models import literals
 from flytekit.models import types as _types
@@ -104,7 +103,7 @@ def test_boolean_primitive():
 
 
 def test_datetime_primitive():
-    dt = datetime.utcnow().replace(tzinfo=pytz.UTC)
+    dt = datetime.utcnow().replace(tzinfo=timezone.utc)
     obj = literals.Primitive(datetime=dt)
     assert obj.integer is None
     assert obj.boolean is None