Skip to content

Commit 7e92e22

Browse files
authored
Merge pull request #8659 from neilflood/jitclass-preserve-modulename
Preserve module name of jitted class
2 parents fe8003d + a40f099 commit 7e92e22

2 files changed

Lines changed: 17 additions & 1 deletion

File tree

‎numba/experimental/jitclass/decorators.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,13 @@ def wrap(cls):
7474
else:
7575
from numba.experimental.jitclass.base import (register_class_type,
7676
ClassBuilder)
77-
return register_class_type(cls, spec, types.ClassType, ClassBuilder)
77+
cls_jitted = register_class_type(cls, spec, types.ClassType,
78+
ClassBuilder)
79+
80+
# Preserve the module name of the original class
81+
cls_jitted.__module__ = cls.__module__
82+
83+
return cls_jitted
7884

7985
if cls_or_spec is None:
8086
return wrap

‎numba/tests/test_jitclasses.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1165,6 +1165,16 @@ def __enter__(self):
11651165
Foo()
11661166
self.assertIn("Method '__enter__' is not supported.", str(e.exception))
11671167

1168+
def test_modulename(self):
1169+
@jitclass
1170+
class TestModname(object):
1171+
def __init__(self):
1172+
self.x = 12
1173+
1174+
thisModule = __name__
1175+
classModule = TestModname.__module__
1176+
self.assertEqual(thisModule, classModule)
1177+
11681178

11691179
class TestJitClassOverloads(MemoryLeakMixin, TestCase):
11701180

0 commit comments

Comments
 (0)