python3.9 parsing different from python3.8

[blythed]

🐞 Bug

This doesn't save in python3.9 although works in python3.8:

import numpy

@transform
def rescale_image_tensor(x):
    return (255 * (x * 0.5 + 0.5)).numpy().astype(numpy.uint8)

Error:

---------------------------------------------------------------------------
NameNotFound                              Traceback (most recent call last)
Input In [4], in <module>
----> 1 rescale_image_tensor.pd_save('test', force_overwrite=True)

File ~/Deeplearning/padl/padl/transforms.py:382, in Transform.pd_save(self, path, force_overwrite)
    380 for i, subtrans in enumerate(self._pd_all_transforms()):
    381     subtrans.pd_pre_save(path, i)
--> 382 code, versions = self._pd_dumps(True, path=path)
    384 with open(path / 'transform.py', 'w') as f:
    385     f.write(code)

File ~/Deeplearning/padl/padl/transforms.py:586, in Transform._pd_dumps(self, return_versions, path)
    578 def _pd_dumps(self, return_versions: bool = False,
    579               path: Optional[Path] = None) -> Union[str, Tuple[str, str]]:
    580     """Dump the transform as python code.
    581 
    582     :param return_versions: If *True* return a tuple of the code and a file listing
    583         dependencies and their versions.
    584     :param path: Optional path to save at, might be required for serializer code snippets.
    585     """
--> 586     graph = self._pd_build_codegraph(name='_pd_main')
    587     Serializer.save_all(graph, path)
    588     code = graph.dumps()

File ~/Deeplearning/padl/padl/transforms.py:510, in Transform._pd_build_codegraph(self, graph, name)
    507     continue
    509 # find how *next_name* came into being
--> 510 next_codenode = find_codenode(next_name)
    511 graph[next_name] = next_codenode
    513 todo.update(next_codenode.globals_)

File ~/Deeplearning/padl/padl/dumptools/var2mod.py:548, in find_codenode(name)
    546 def find_codenode(name: ScopedName):
    547     """Find the :class:`CodeNode` corresponding to a :class:`ScopedName` *name*. """
--> 548     (source, node), scope_of_next_var = find_in_scope(name)
    550     # find dependencies
    551     globals_ = find_globals(node)

File ~/Deeplearning/padl/padl/dumptools/symfinder.py:636, in find_in_scope(name)
    634 if scope.module is None:
    635     raise NameNotFound(f'{name.name} not found in function hierarchy.')
--> 636 source, node = find(name.name, scope.module, i)
    637 if getattr(node, '_globalscope', False):
    638     scope = Scope.empty()

File ~/Deeplearning/padl/padl/dumptools/symfinder.py:768, in find(var_name, module, i)
    766     module = sys.modules['__main__']
    767 try:
--> 768     return find_in_module(var_name, module, i)
    769 except TypeError as exc:
    770     if module is not sys.modules['__main__']:

File ~/Deeplearning/padl/padl/dumptools/symfinder.py:700, in find_in_module(var_name, module, i)
    692 """Find the piece of code that assigned a value to the variable with name *var_name* in the
    693 module *module*.
    694 
   (...)
    697 :returns: Tuple with source code segment and corresponding ast node.
    698 """
    699 source = sourceget.get_module_source(module)
--> 700 return find_in_source(var_name, source, i=i)

File ~/Deeplearning/padl/padl/dumptools/symfinder.py:688, in find_in_source(var_name, source, tree, i, return_partial)
    686 if return_partial:
    687     return i
--> 688 raise NameNotFound(f'{var_name} not found.')

NameNotFound: x not found.