test_source.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. #!/usr/bin/env python
  2. #
  3. # Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
  4. # Copyright (c) 2008-2016 California Institute of Technology.
  5. # Copyright (c) 2016-2024 The Uncertainty Quantification Foundation.
  6. # License: 3-clause BSD. The full license text is available at:
  7. # - https://github.com/uqfoundation/dill/blob/master/LICENSE
  8. from dill.source import getsource, getname, _wrap, getimport
  9. from dill.source import importable
  10. from dill._dill import IS_PYPY
  11. import sys
  12. PY310b = 0x30a00b1
  13. f = lambda x: x**2
  14. def g(x): return f(x) - x
  15. def h(x):
  16. def g(x): return x
  17. return g(x) - x
  18. class Foo(object):
  19. def bar(self, x):
  20. return x*x+x
  21. _foo = Foo()
  22. def add(x,y):
  23. return x+y
  24. # yes, same as 'f', but things are tricky when it comes to pointers
  25. squared = lambda x:x**2
  26. class Bar:
  27. pass
  28. _bar = Bar()
  29. # inspect.getsourcelines # dill.source.getblocks
  30. def test_getsource():
  31. assert getsource(f) == 'f = lambda x: x**2\n'
  32. assert getsource(g) == 'def g(x): return f(x) - x\n'
  33. assert getsource(h) == 'def h(x):\n def g(x): return x\n return g(x) - x\n'
  34. assert getname(f) == 'f'
  35. assert getname(g) == 'g'
  36. assert getname(h) == 'h'
  37. assert _wrap(f)(4) == 16
  38. assert _wrap(g)(4) == 12
  39. assert _wrap(h)(4) == 0
  40. assert getname(Foo) == 'Foo'
  41. assert getname(Bar) == 'Bar'
  42. assert getsource(Bar) == 'class Bar:\n pass\n'
  43. assert getsource(Foo) == 'class Foo(object):\n def bar(self, x):\n return x*x+x\n'
  44. #XXX: add getsource for _foo, _bar
  45. # test itself
  46. def test_itself():
  47. assert getimport(getimport)=='from dill.source import getimport\n'
  48. # builtin functions and objects
  49. def test_builtin():
  50. assert getimport(pow) == 'pow\n'
  51. assert getimport(100) == '100\n'
  52. assert getimport(True) == 'True\n'
  53. assert getimport(pow, builtin=True) == 'from builtins import pow\n'
  54. assert getimport(100, builtin=True) == '100\n'
  55. assert getimport(True, builtin=True) == 'True\n'
  56. # this is kinda BS... you can't import a None
  57. assert getimport(None) == 'None\n'
  58. assert getimport(None, builtin=True) == 'None\n'
  59. # other imported functions
  60. def test_imported():
  61. from math import sin
  62. assert getimport(sin) == 'from math import sin\n'
  63. # interactively defined functions
  64. def test_dynamic():
  65. assert getimport(add) == 'from %s import add\n' % __name__
  66. # interactive lambdas
  67. assert getimport(squared) == 'from %s import squared\n' % __name__
  68. # classes and class instances
  69. def test_classes():
  70. from io import BytesIO as StringIO
  71. y = "from _io import BytesIO\n"
  72. x = y if (IS_PYPY or sys.hexversion >= PY310b) else "from io import BytesIO\n"
  73. s = StringIO()
  74. assert getimport(StringIO) == x
  75. assert getimport(s) == y
  76. # interactively defined classes and class instances
  77. assert getimport(Foo) == 'from %s import Foo\n' % __name__
  78. assert getimport(_foo) == 'from %s import Foo\n' % __name__
  79. # test importable
  80. def test_importable():
  81. assert importable(add, source=False) == 'from %s import add\n' % __name__
  82. assert importable(squared, source=False) == 'from %s import squared\n' % __name__
  83. assert importable(Foo, source=False) == 'from %s import Foo\n' % __name__
  84. assert importable(Foo.bar, source=False) == 'from %s import bar\n' % __name__
  85. assert importable(_foo.bar, source=False) == 'from %s import bar\n' % __name__
  86. assert importable(None, source=False) == 'None\n'
  87. assert importable(100, source=False) == '100\n'
  88. assert importable(add, source=True) == 'def add(x,y):\n return x+y\n'
  89. assert importable(squared, source=True) == 'squared = lambda x:x**2\n'
  90. assert importable(None, source=True) == 'None\n'
  91. assert importable(Bar, source=True) == 'class Bar:\n pass\n'
  92. assert importable(Foo, source=True) == 'class Foo(object):\n def bar(self, x):\n return x*x+x\n'
  93. assert importable(Foo.bar, source=True) == 'def bar(self, x):\n return x*x+x\n'
  94. assert importable(Foo.bar, source=False) == 'from %s import bar\n' % __name__
  95. assert importable(Foo.bar, alias='memo', source=False) == 'from %s import bar as memo\n' % __name__
  96. assert importable(Foo, alias='memo', source=False) == 'from %s import Foo as memo\n' % __name__
  97. assert importable(squared, alias='memo', source=False) == 'from %s import squared as memo\n' % __name__
  98. assert importable(squared, alias='memo', source=True) == 'memo = squared = lambda x:x**2\n'
  99. assert importable(add, alias='memo', source=True) == 'def add(x,y):\n return x+y\n\nmemo = add\n'
  100. assert importable(None, alias='memo', source=True) == 'memo = None\n'
  101. assert importable(100, alias='memo', source=True) == 'memo = 100\n'
  102. assert importable(add, builtin=True, source=False) == 'from %s import add\n' % __name__
  103. assert importable(squared, builtin=True, source=False) == 'from %s import squared\n' % __name__
  104. assert importable(Foo, builtin=True, source=False) == 'from %s import Foo\n' % __name__
  105. assert importable(Foo.bar, builtin=True, source=False) == 'from %s import bar\n' % __name__
  106. assert importable(_foo.bar, builtin=True, source=False) == 'from %s import bar\n' % __name__
  107. assert importable(None, builtin=True, source=False) == 'None\n'
  108. assert importable(100, builtin=True, source=False) == '100\n'
  109. def test_numpy():
  110. try:
  111. import numpy as np
  112. y = np.array
  113. x = y([1,2,3])
  114. assert importable(x, source=False) == 'from numpy import array\narray([1, 2, 3])\n'
  115. assert importable(y, source=False) == 'from %s import array\n' % y.__module__
  116. assert importable(x, source=True) == 'from numpy import array\narray([1, 2, 3])\n'
  117. assert importable(y, source=True) == 'from %s import array\n' % y.__module__
  118. y = np.int64
  119. x = y(0)
  120. assert importable(x, source=False) == 'from numpy import int64\nint64(0)\n'
  121. assert importable(y, source=False) == 'from %s import int64\n' % y.__module__
  122. assert importable(x, source=True) == 'from numpy import int64\nint64(0)\n'
  123. assert importable(y, source=True) == 'from %s import int64\n' % y.__module__
  124. y = np.bool_
  125. x = y(0)
  126. import warnings
  127. with warnings.catch_warnings():
  128. warnings.filterwarnings('ignore', category=FutureWarning)
  129. warnings.filterwarnings('ignore', category=DeprecationWarning)
  130. if hasattr(np, 'bool'): b = 'bool_' if np.bool is bool else 'bool'
  131. else: b = 'bool_'
  132. assert importable(x, source=False) == 'from numpy import %s\n%s(False)\n' % (b,b)
  133. assert importable(y, source=False) == 'from %s import %s\n' % (y.__module__,b)
  134. assert importable(x, source=True) == 'from numpy import %s\n%s(False)\n' % (b,b)
  135. assert importable(y, source=True) == 'from %s import %s\n' % (y.__module__,b)
  136. except ImportError: pass
  137. #NOTE: if before getimport(pow), will cause pow to throw AssertionError
  138. def test_foo():
  139. assert importable(_foo, source=True).startswith("import dill\nclass Foo(object):\n def bar(self, x):\n return x*x+x\ndill.loads(")
  140. if __name__ == '__main__':
  141. test_getsource()
  142. test_itself()
  143. test_builtin()
  144. test_imported()
  145. test_dynamic()
  146. test_classes()
  147. test_importable()
  148. test_numpy()
  149. test_foo()