Source code for objax.util.tracing

# Copyright 2021 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

__all__ = ['find_used_variables']

import inspect
import ast
import re
import tokenize

from io import StringIO

from typing import Callable
from objax.variable import BaseVar, VarCollection
from objax.module import Module

def getanno(node, key, field_name='__anno'):
    """Gets annotation of the AST node."""
    return getattr(node, field_name)[key]

def hasanno(node, key, field_name='__anno'):
    """Returns whether AST node has an annotation."""
    return hasattr(node, field_name) and key in getattr(node, field_name)

def setanno(node, key, value, field_name='__anno'):
    """Sets annotation on AST node."""
    annotations = getattr(node, field_name, {})
    setattr(node, field_name, annotations)
    annotations[key] = value

class AnalyzeUserVariablesNodeTransformer(ast.NodeTransformer):

    def __init__(self, closure_vars, global_vars):
        self.closure_vars = closure_vars
        self.global_vars = global_vars = VarCollection()

    def check_objax_var_module(self, node):
        if not hasanno(node, 'value'):
        v = getanno(node, 'value')
        v_name = getanno(node, 'name')
        if v is None:
        if isinstance(v, Module):
   + '.'))
            setanno(node, 'value', None)
        if isinstance(v, BaseVar):
            if v_name in and[v_name] is not v:
                # This generally should not happen and probably indication of a bug somewhere.
                raise ValueError(
                    f'Variable tracing failed because two variables were found with the same name {v_name}')
      [v_name] = v
                setanno(node, 'value', None)

    def visit_Name(self, node):
        node = self.generic_visit(node)
        if isinstance(node.ctx, ast.Load):
            if in self.closure_vars:
                setanno(node, 'name',
                setanno(node, 'value', self.closure_vars[])
            elif in self.global_vars:
                setanno(node, 'name',
                setanno(node, 'value', self.global_vars[])
        return node

    def visit_Attribute(self, node):
        node = self.generic_visit(node)
        if isinstance(node.ctx, ast.Load) and hasanno(node.value, 'value'):
            parent_value = getanno(node.value, 'value')
            if parent_value is not None and hasattr(parent_value, node.attr):
                setanno(node, 'name', getanno(node.value, 'name') + '.' + node.attr)
                setanno(node, 'value', getattr(parent_value, node.attr))

        return node

_LEADING_WHITESPACE = re.compile(r'\s*')

def dedent_block(code_string):
    """Dedents a code so that its first line starts at row zero."""

    # Removes any backslash line continuations from the code
    code_string = code_string.replace('\\\n', '')

    token_gen = tokenize.generate_tokens(StringIO(code_string).readline)

    block_indentation = None
    tokens = []
        for tok in token_gen:
    except tokenize.TokenError:
        # Resolution of lambda functions may yield incomplete code, which can
        # in turn generate this error. We silently ignore this error because the
        # parser may still be able to deal with it.

    for tok in tokens:
        tok_type, tok_string, _, _, _ = tok
        if tok_type == tokenize.INDENT:
            block_indentation = tok_string
        elif tok_type not in (tokenize.NL, tokenize.NEWLINE, tokenize.STRING, tokenize.COMMENT):
            block_indentation = ''

    if not block_indentation:
        return code_string

    block_level = len(block_indentation)
    first_indent_uses_tabs = '\t' in block_indentation
    for i, tok in enumerate(tokens):
        tok_type, tok_string, _, _, _ = tok
        if tok_type == tokenize.INDENT:
            if ((' ' in tok_string and first_indent_uses_tabs) or ('\t' in tok_string and not first_indent_uses_tabs)):
                raise ValueError('Code mixing tabs and spaces for indentation is not allowed')
            if len(tok_string) >= block_level:
                tok_string = tok_string[block_level:]
            tokens[i] = (tok_type, tok_string)

    new_code = tokenize.untokenize(tokens)

    # Note: untokenize respects the line structure, but not the whitespace within
    # lines. For example, `def foo()` may be untokenized as `def foo ()`
    # So instead of using the output of dedent, we match the leading whitespace
    # on each line.
    dedented_code = []
    for line, new_line in zip(code_string.split('\n'), new_code.split('\n')):
        original_indent = re.match(_LEADING_WHITESPACE, line).group()
        new_indent = re.match(_LEADING_WHITESPACE, new_line).group()
        if len(original_indent) > len(new_indent):
            dedented_line = line[len(original_indent) - len(new_indent):]
            dedented_line = line
    new_code = '\n'.join(dedented_code)

    return new_code

[docs]def find_used_variables(fn: Callable) -> VarCollection: """Finds all Objax variables which are used by a given callable. Args: fn: input function or callable. Returns: Variable collection with all variables used by input function. """ if not hasattr(fn, '__code__'): raise ValueError('Can not determine variables used by a function. Function does not have __code__ attribute.') try: src = inspect.getsource(fn) except OSError: raise ValueError('Can not determine variables used by a function. Code of the function can not be retrieved.') src = dedent_block(src) main_node = ast.parse(src) main_node = main_node.body[0] if fn.__closure__: closure_vars = {name: cell.cell_contents for name, cell in zip(fn.__code__.co_freevars, fn.__closure__)} else: closure_vars = {} analyzer = AnalyzeUserVariablesNodeTransformer(closure_vars=closure_vars, global_vars=fn.__globals__) analyzer.visit(main_node) return