match on both

pola-rs · Aug 5, 2023 · 73ee469 · 73ee469
1 parent 68efb5f
commit 73ee469
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/py-polars/polars/utils/udfs.py b/py-polars/polars/utils/udfs.py
@@ -34,12 +34,6 @@ class StackValue(NamedTuple):
 StackEntry: TypeAlias = Union[str, StackValue]
 
 _MIN_PY311 = sys.version_info >= (3, 11)
-try:
-    get_ipython  # type: ignore[name-defined]  # noqa: B018
-except NameError:
-    _IN_IPYTHON = False
-else:
-    _IN_IPYTHON = True
 
 
 class OpNames:
@@ -72,7 +66,7 @@ class OpNames:
         }
     )
     LOAD_VALUES = frozenset(("LOAD_CONST", "LOAD_DEREF", "LOAD_FAST", "LOAD_GLOBAL"))
-    LOAD_ATTR = {"LOAD_ATTR"} if (_MIN_PY311 and not _IN_IPYTHON) else {"LOAD_METHOD"}
+    LOAD_ATTR = {"LOAD_METHOD", "LOAD_ATTR"} if _MIN_PY311 else {"LOAD_METHOD"}
     LOAD = LOAD_VALUES | {"LOAD_METHOD", "LOAD_ATTR"}
     SYNTHETIC = {
         "POLARS_EXPRESSION": 1,

diff --git a/py-polars/tests/test_udfs.py b/py-polars/tests/test_udfs.py
@@ -176,6 +176,8 @@ def test_bytecode_parser_expression_in_ipython(
 
     script = (
         "import udfs; "
+        "import numpy as np; "
+        "import json; "
         f"MY_DICT = {MY_DICT};"
         f'bytecode_parser = udfs.BytecodeParser({func}, apply_target="expr");'
         f'print(bytecode_parser.to_expression("{col}"));'