Skip to content

Commit 5197876

Browse files
authored
Merge pull request #6999 from sklam/fix/cloudpickle_fallout
Fix errors introduced by the cloudpickle patch
2 parents 2772b84 + a7267d2 commit 5197876

3 files changed

Lines changed: 50 additions & 10 deletions

File tree

‎numba/core/serialize.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ def dumps(obj):
5252
"""
5353
pickler = NumbaPickler
5454
with io.BytesIO() as buf:
55-
p = pickler(buf)
55+
p = pickler(buf, protocol=4)
5656
p.dump(obj)
5757
pickled = buf.getvalue()
5858

‎numba/core/utils.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,15 @@
66
import operator
77
import timeit
88
import math
9+
import sys
910
import traceback
1011
import weakref
1112
import warnings
1213
import threading
1314
import contextlib
1415

1516
from types import ModuleType
17+
from importlib import import_module
1618
from collections.abc import Mapping, Sequence
1719
import numpy as np
1820

@@ -537,3 +539,47 @@ def unified_function_type(numba_types, require_precise=True):
537539
function = types.UndefinedFunctionType(mnargs, dispatchers)
538540

539541
return function
542+
543+
544+
class _RedirectSubpackage(ModuleType):
545+
"""Redirect a subpackage to a subpackage.
546+
547+
This allows all references like:
548+
549+
>>> from numba.old_subpackage import module
550+
>>> module.item
551+
552+
>>> import numba.old_subpackage.module
553+
>>> numba.old_subpackage.module.item
554+
555+
>>> from numba.old_subpackage.module import item
556+
"""
557+
def __init__(self, old_module_locals, new_module):
558+
old_module = old_module_locals['__name__']
559+
super().__init__(old_module)
560+
561+
self.__old_module_states = {}
562+
self.__new_module = new_module
563+
564+
new_mod_obj = import_module(new_module)
565+
566+
# Map all sub-modules over
567+
for k, v in new_mod_obj.__dict__.items():
568+
# Get attributes so that `subpackage.xyz` and
569+
# `from subpackage import xyz` work
570+
setattr(self, k, v)
571+
if isinstance(v, ModuleType):
572+
# Map modules into the interpreter so that
573+
# `import subpackage.xyz` works
574+
sys.modules[f"{old_module}.{k}"] = sys.modules[v.__name__]
575+
576+
# copy across dunders so that package imports work too
577+
for attr, value in old_module_locals.items():
578+
if attr.startswith('__') and attr.endswith('__'):
579+
if attr != "__builtins__":
580+
setattr(self, attr, value)
581+
self.__old_module_states[attr] = value
582+
583+
def __reduce__(self):
584+
args = (self.__old_module_states, self.__new_module)
585+
return _RedirectSubpackage, args

‎numba/types/__init__.py‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,3 @@
1-
import importlib
2-
3-
4-
def __getattr__(name):
5-
"""Redirect to numba.core.types
6-
"""
7-
# Uses PEP-562
8-
basemod = importlib.import_module(f"numba.core.types", __name__)
9-
return getattr(basemod, name)
1+
import sys
2+
from numba.core.utils import _RedirectSubpackage
3+
sys.modules[__name__] = _RedirectSubpackage(locals(), "numba.core.types")

0 commit comments

Comments
 (0)