From 48eff98f51d52c32d20a680d2115da8e2c09828c Mon Sep 17 00:00:00 2001
From: Sammy Sidhu <samster25@users.noreply.github.com>
Date: Mon, 15 Apr 2024 12:47:55 -0700
Subject: [PATCH] relax linting on cloudpickling

---
 .ruff.toml                 |  2 ++
 daft/pickle/cloudpickle.py | 50 +++++++++++++++++++++++++++++++++++---
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/.ruff.toml b/.ruff.toml
index 5c1627635d..18a03ebabb 100644
--- a/.ruff.toml
+++ b/.ruff.toml
@@ -29,5 +29,7 @@ ignore = [
 [lint.per-file-ignores]
 # Do not enforce usage and import order rules in init files
 "__init__.py" = ["E402", "F401", "I"]
+# ignore vendored cloudpickle old python code
+"daft/pickle/cloudpickle*" = ["UP036"]
 # Allow wild imports in conftest
 "tests/conftest.py" = ["F405", "E402", "F403"]
diff --git a/daft/pickle/cloudpickle.py b/daft/pickle/cloudpickle.py
index 0c68c7f894..4e49deb51c 100644
--- a/daft/pickle/cloudpickle.py
+++ b/daft/pickle/cloudpickle.py
@@ -437,7 +437,17 @@ def cell_set(cell, value):
     test and checker libraries decide to parse the whole file.
     """
 
-    cell.cell_contents = value
+    if sys.version_info[:2] >= (3, 7):  # pragma: no branch
+        cell.cell_contents = value
+    else:
+        _cell_set = types.FunctionType(
+            _cell_set_template_code,
+            {},
+            "_cell_set",
+            (),
+            (cell,),
+        )
+        _cell_set(value)
 
 
 def _make_cell_set_template_code():
@@ -467,6 +477,9 @@ def _cell_set_factory(value):
     return _cell_set_template_code
 
 
+if sys.version_info[:2] < (3, 7):
+    _cell_set_template_code = _make_cell_set_template_code()
+
 # relevant opcodes
 STORE_GLOBAL = opcode.opmap["STORE_GLOBAL"]
 DELETE_GLOBAL = opcode.opmap["DELETE_GLOBAL"]
@@ -523,8 +536,39 @@ def _extract_class_dict(cls):
     return clsdict
 
 
-_is_parametrized_type_hint = None
-_create_parametrized_type_hint = None
+if sys.version_info[:2] < (3, 7):  # pragma: no branch
+
+    def _is_parametrized_type_hint(obj):
+        # This is very cheap but might generate false positives. So try to
+        # narrow it down is good as possible.
+        type_module = getattr(type(obj), "__module__", None)
+        from_typing_extensions = type_module == "typing_extensions"
+        from_typing = type_module == "typing"
+
+        # general typing Constructs
+        is_typing = getattr(obj, "__origin__", None) is not None
+
+        # typing_extensions.Literal
+        is_literal = (getattr(obj, "__values__", None) is not None) and from_typing_extensions
+
+        # typing_extensions.Final
+        is_final = (getattr(obj, "__type__", None) is not None) and from_typing_extensions
+
+        # typing.ClassVar
+        is_classvar = (getattr(obj, "__type__", None) is not None) and from_typing
+
+        # typing.Union/Tuple for old Python 3.5
+        is_union = getattr(obj, "__union_params__", None) is not None
+        is_tuple = getattr(obj, "__tuple_params__", None) is not None
+        is_callable = getattr(obj, "__result__", None) is not None and getattr(obj, "__args__", None) is not None
+        return any((is_typing, is_literal, is_final, is_classvar, is_union, is_tuple, is_callable))
+
+    def _create_parametrized_type_hint(origin, args):
+        return origin[args]
+
+else:
+    _is_parametrized_type_hint = None
+    _create_parametrized_type_hint = None
 
 
 def parametrized_type_hint_getinitargs(obj):