test_source.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  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-2023 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, likely_import
  9. from dill.source import getimportable
  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 likely_import(likely_import)=='from dill.source import likely_import\n'
  48. # builtin functions and objects
  49. def test_builtin():
  50. assert likely_import(pow) == 'pow\n'
  51. assert likely_import(100) == '100\n'
  52. assert likely_import(True) == 'True\n'
  53. assert likely_import(pow, explicit=True) == 'from builtins import pow\n'
  54. assert likely_import(100, explicit=True) == '100\n'
  55. assert likely_import(True, explicit=True) == 'True\n'
  56. # this is kinda BS... you can't import a None
  57. assert likely_import(None) == 'None\n'
  58. assert likely_import(None, explicit=True) == 'None\n'
  59. # other imported functions
  60. def test_imported():
  61. from math import sin
  62. assert likely_import(sin) == 'from math import sin\n'
  63. # interactively defined functions
  64. def test_dynamic():
  65. assert likely_import(add) == 'from %s import add\n' % __name__
  66. # interactive lambdas
  67. assert likely_import(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 likely_import(StringIO) == x
  75. assert likely_import(s) == y
  76. # interactively defined classes and class instances
  77. assert likely_import(Foo) == 'from %s import Foo\n' % __name__
  78. assert likely_import(_foo) == 'from %s import Foo\n' % __name__
  79. # test getimportable
  80. def test_importable():
  81. assert getimportable(add) == 'from %s import add\n' % __name__
  82. assert getimportable(squared) == 'from %s import squared\n' % __name__
  83. assert getimportable(Foo) == 'from %s import Foo\n' % __name__
  84. assert getimportable(Foo.bar) == 'from %s import bar\n' % __name__
  85. assert getimportable(_foo.bar) == 'from %s import bar\n' % __name__
  86. assert getimportable(None) == 'None\n'
  87. assert getimportable(100) == '100\n'
  88. assert getimportable(add, byname=False) == 'def add(x,y):\n return x+y\n'
  89. assert getimportable(squared, byname=False) == 'squared = lambda x:x**2\n'
  90. assert getimportable(None, byname=False) == 'None\n'
  91. assert getimportable(Bar, byname=False) == 'class Bar:\n pass\n'
  92. assert getimportable(Foo, byname=False) == 'class Foo(object):\n def bar(self, x):\n return x*x+x\n'
  93. assert getimportable(Foo.bar, byname=False) == 'def bar(self, x):\n return x*x+x\n'
  94. assert getimportable(Foo.bar, byname=True) == 'from %s import bar\n' % __name__
  95. assert getimportable(Foo.bar, alias='memo', byname=True) == 'from %s import bar as memo\n' % __name__
  96. assert getimportable(Foo, alias='memo', byname=True) == 'from %s import Foo as memo\n' % __name__
  97. assert getimportable(squared, alias='memo', byname=True) == 'from %s import squared as memo\n' % __name__
  98. assert getimportable(squared, alias='memo', byname=False) == 'memo = squared = lambda x:x**2\n'
  99. assert getimportable(add, alias='memo', byname=False) == 'def add(x,y):\n return x+y\n\nmemo = add\n'
  100. assert getimportable(None, alias='memo', byname=False) == 'memo = None\n'
  101. assert getimportable(100, alias='memo', byname=False) == 'memo = 100\n'
  102. assert getimportable(add, explicit=True) == 'from %s import add\n' % __name__
  103. assert getimportable(squared, explicit=True) == 'from %s import squared\n' % __name__
  104. assert getimportable(Foo, explicit=True) == 'from %s import Foo\n' % __name__
  105. assert getimportable(Foo.bar, explicit=True) == 'from %s import bar\n' % __name__
  106. assert getimportable(_foo.bar, explicit=True) == 'from %s import bar\n' % __name__
  107. assert getimportable(None, explicit=True) == 'None\n'
  108. assert getimportable(100, explicit=True) == '100\n'
  109. def test_numpy():
  110. try:
  111. from numpy import array
  112. x = array([1,2,3])
  113. assert getimportable(x) == 'from numpy import array\narray([1, 2, 3])\n'
  114. assert getimportable(array) == 'from %s import array\n' % array.__module__
  115. assert getimportable(x, byname=False) == 'from numpy import array\narray([1, 2, 3])\n'
  116. assert getimportable(array, byname=False) == 'from %s import array\n' % array.__module__
  117. except ImportError: pass
  118. #NOTE: if before likely_import(pow), will cause pow to throw AssertionError
  119. def test_foo():
  120. assert getimportable(_foo, byname=False).startswith("import dill\nclass Foo(object):\n def bar(self, x):\n return x*x+x\ndill.loads(")
  121. if __name__ == '__main__':
  122. test_getsource()
  123. test_itself()
  124. test_builtin()
  125. test_imported()
  126. test_dynamic()
  127. test_classes()
  128. test_importable()
  129. test_numpy()
  130. test_foo()