Python ast.ImportFrom() Examples

def _find_imports(self, node):
        """Recurses through AST collecting the targets of all import
        statements."""
        if isinstance(node, ast.Import):
            return {self._extract_root_module(alias.name) for alias in node.names}
        elif isinstance(node, ast.ImportFrom):
            # We ignore all imports with levels other than 0. That's because if
            # if level > 0, we know that it's a relative import, and we only
            # care about root modules.
            if node.level == 0:
                return {self._extract_root_module(node.module)}
            else:
                return set()
        elif hasattr(node, 'body') and hasattr(node.body, '__iter__'):
            # Not all bodies are lists (for ex. exec)
            imps = set()
            for child_node in node.body:
                imps.update(self._find_imports(child_node))
            return imps
        else:
            return set() 

def _find_imports(self, node):
        """Recurses through AST collecting the targets of all import
        statements."""
        if isinstance(node, ast.Import):
            return {self._extract_root_module(alias.name) for alias in node.names}
        elif isinstance(node, ast.ImportFrom):
            # We ignore all imports with levels other than 0. That's because if
            # if level > 0, we know that it's a relative import, and we only
            # care about root modules.
            if node.level == 0:
                return {self._extract_root_module(node.module)}
            else:
                return set()
        elif hasattr(node, 'body') and hasattr(node.body, '__iter__'):
            # Not all bodies are lists (for ex. exec)
            imps = set()
            for child_node in node.body:
                imps.update(self._find_imports(child_node))
            return imps
        else:
            return set() 

def handleNode(self, node, parent):
        if node is None:
            return
        if self.offset and getattr(node, 'lineno', None) is not None:
            node.lineno += self.offset[0]
            node.col_offset += self.offset[1]
        if self.traceTree:
            print('  ' * self.nodeDepth + node.__class__.__name__)
        if self.futuresAllowed and not (isinstance(node, ast.ImportFrom) or
                                        self.isDocstring(node)):
            self.futuresAllowed = False
        self.nodeDepth += 1
        node.depth = self.nodeDepth
        node.parent = parent
        try:
            handler = self.getNodeHandler(node.__class__)
            handler(node)
        finally:
            self.nodeDepth -= 1
        if self.traceTree:
            print('  ' * self.nodeDepth + 'end ' + node.__class__.__name__) 

def _parse_mock_imports(mod_ast, expanded_imports):
  """Parses a module AST node for import statements and resolves them against
  expanded_imports (such as you might get from _expand_mock_imports).

  If an import is not recognized, it is omitted from the returned dictionary.

  Returns a dictionary suitable for eval'ing a statement in mod_ast, with
  symbols from mod_ast's imports resolved to real objects, as per
  expanded_imports.
  """
  ret = {}

  for node in mod_ast.body:
    if isinstance(node, ast.Import):
      for alias in node.names:
        if alias.name in expanded_imports:
          ret[alias.asname or alias.name] = expanded_imports[alias.name]
    elif isinstance(node, ast.ImportFrom):
      if node.level == 0:
        for alias in node.names:
          fullname ='%s.%s' % (node.module, alias.name)
          if fullname in expanded_imports:
            ret[alias.asname or alias.name] = expanded_imports[fullname]

  return ret 

def _get_modules(self, content) -> Set[str]:
        imports = set()
        tree = ast.parse(content)
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for subnode in node.names:
                    imports.add(subnode.name)
            elif isinstance(node, ast.ImportFrom) and node.level == 0:
                imports.add(node.module)
        modules = set()
        for module in imports:
            if not module:
                continue
            module = module.split('.', maxsplit=1)[0]
            if module in self.stdlib:
                continue
            module = self.aliases.get(module, module)
            modules.add(module)
        return modules 

def get_top_imported_names(file: str) -> Set[str]:
    """Collect names imported in given file.

    We only collect top-level names, i.e. `from foo.bar import baz`
    will only add `foo` to the list.
    """
    if not file.endswith(".pyi"):
        return set()
    with open(os.path.join(file), "rb") as f:
        content = f.read()
    parsed = ast.parse(content)
    top_imported = set()
    for node in ast.walk(parsed):
        if isinstance(node, ast.Import):
            for name in node.names:
                top_imported.add(name.name.split('.')[0])
        elif isinstance(node, ast.ImportFrom):
            if node.level > 0:
                # Relative imports always refer to the current package.
                continue
            assert node.module
            top_imported.add(node.module.split('.')[0])
    return top_imported 

def visit(self, node) -> Any:
        if isinstance(node, ast.Name):
            if isinstance(node.ctx, ast.Load):
                self.loaded.add(node.id)
            elif isinstance(node.ctx, ast.Store):
                self.stored.add(node.id)
        elif isinstance(node, ast.Return):
            self.has_return = True
        # We must keep track of importer name in order to avoid considering as variable
        # names:
        elif isinstance(node, ast.Import):
            self.imported.update([ n.name for n in node.names])
        elif isinstance(node, ast.ImportFrom):
            self.imported.update([ n.name for n in node.names])

        self.generic_visit(node) 

def add_custom_type_symbol(self,
                               custom_type: ir2.CustomType,
                               definition_ast_node: Union[ast.ClassDef, ast.ImportFrom, None],
                               source_module: Optional[str] = None):
        self.add_symbol(name=custom_type.name,
                        expr_type=ir2.FunctionType(argtypes=tuple(arg.expr_type
                                                                  for arg in custom_type.arg_types),
                                                   argnames=tuple(arg.name
                                                                  for arg in custom_type.arg_types),
                                                   returns=custom_type),
                        definition_ast_node=definition_ast_node,
                        is_only_partially_defined=False,
                        is_function_that_may_throw=False)
        self.custom_types_symbol_table.add_symbol(name=custom_type.name,
                                                  expr_type=custom_type,
                                                  definition_ast_node=definition_ast_node,
                                                  is_only_partially_defined=False,
                                                  is_function_that_may_throw=False,
                                                  source_module=source_module) 

def add_symbol_for_external_elem(self,
                                     elem: Union[ir2.FunctionDefn, ir2.CustomType],
                                     import_from_ast_node: ast.ImportFrom,
                                     source_module: str):
        if isinstance(elem, ir2.FunctionDefn):
            self.add_symbol(name=elem.name,
                            expr_type=ir2.FunctionType(argtypes=tuple(arg.expr_type for arg in elem.args),
                                                       argnames=tuple(arg.name for arg in elem.args),
                                                       returns=elem.return_type),
                            definition_ast_node=import_from_ast_node,
                            is_only_partially_defined=False,
                            is_function_that_may_throw=True,
                            source_module=source_module)
        elif isinstance(elem, ir2.CustomType):
            self.add_custom_type_symbol(custom_type=elem,
                                        definition_ast_node=import_from_ast_node,
                                        source_module=source_module)
        else:
            raise NotImplementedError('Unexpected elem type: %s' % elem.__class__.__name__) 

def find_globals(g, tree):
    """Uses AST to find globals in an ast tree"""
    for child in tree:
        if hasattr(child, 'body') and isinstance(child.body, list):
            find_globals(g, child.body)
        elif isinstance(child, (ast.FunctionDef, ast.ClassDef)):
            g.add(child.name)
            continue
        elif isinstance(child, ast.Assign):
            try:
                g.add(child.targets[0].id)
            except (IndexError, AttributeError):
                pass
        elif isinstance(child, ast.Import):
            g.add(child.names[0].name)
        elif isinstance(child, ast.ImportFrom):
            for name in child.names:
                g_name = name.asname or name.name
                if g_name == '*':
                    continue
                g.add(g_name) 

def get_local_module_reqs(mod):
    tree = ast.parse(inspect.getsource(mod))
    imports = []
    for statement in tree.body:
        if isinstance(statement, ast.Import):
            imports += [(n.name, None) for n in statement.names]
        elif isinstance(statement, ast.ImportFrom):
            if statement.level == 0:
                imp = (statement.module, None)
            else:
                imp = ('.' + statement.module, mod.__package__)
            imports.append(imp)

    result = [import_module(i, p) for i, p in imports]
    if mod.__file__.endswith('__init__.py'):
        # add loaded subpackages
        prefix = mod.__name__ + '.'
        result += [mod for name, mod in sys.modules.items() if name.startswith(prefix)]
    return result 

def handleNode(self, node, parent):
        if node is None:
            return
        if self.offset and getattr(node, 'lineno', None) is not None:
            node.lineno += self.offset[0]
            node.col_offset += self.offset[1]
        if self.traceTree:
            print('  ' * self.nodeDepth + node.__class__.__name__)
        if self.futuresAllowed and not (isinstance(node, ast.ImportFrom) or
                                        self.isDocstring(node)):
            self.futuresAllowed = False
        self.nodeDepth += 1
        node._pyflakes_depth = self.nodeDepth
        node._pyflakes_parent = parent
        try:
            handler = self.getNodeHandler(node.__class__)
            handler(node)
        finally:
            self.nodeDepth -= 1
        if self.traceTree:
            print('  ' * self.nodeDepth + 'end ' + node.__class__.__name__) 

def handleNode(self, node, parent):
        if node is None:
            return
        if self.offset and getattr(node, 'lineno', None) is not None:
            node.lineno += self.offset[0]
            node.col_offset += self.offset[1]
        if self.traceTree:
            print('  ' * self.nodeDepth + node.__class__.__name__)
        if self.futuresAllowed and not (isinstance(node, ast.ImportFrom) or
                                        self.isDocstring(node)):
            self.futuresAllowed = False
        self.nodeDepth += 1
        node.depth = self.nodeDepth
        node.parent = parent
        try:
            handler = self.getNodeHandler(node.__class__)
            handler(node)
        finally:
            self.nodeDepth -= 1
        if self.traceTree:
            print('  ' * self.nodeDepth + 'end ' + node.__class__.__name__) 

def should_trace(source):
    trace_stmt = None
    deep = False
    for stmt in ast.parse(source).body:
        if isinstance(stmt, ast.Import):
            for alias in stmt.names:
                if alias.name.startswith('birdseye.trace_module'):
                    trace_stmt = stmt
                    if alias.name.endswith('deep'):
                        deep = True

        if isinstance(stmt, ast.ImportFrom) and stmt.module == 'birdseye':
            for alias in stmt.names:
                if alias.name.startswith('trace_module'):
                    trace_stmt = stmt
                    if alias.name.endswith('deep'):
                        deep = True
    return deep, trace_stmt 

def find_imports(code):
    """
    Finds the imports in a string of code and returns a list of their package
    names.
    """
    # handle mis-indented input from multi-line strings
    code = dedent(code)

    mod = ast.parse(code)
    imports = set()
    for node in ast.walk(mod):
        if isinstance(node, ast.Import):
            for name in node.names:
                name = name.name
                imports.add(name.split(".")[0])
        elif isinstance(node, ast.ImportFrom):
            name = node.module
            imports.add(name.split(".")[0])
    return list(imports) 

def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        self.imports.add(_to_offset(node))
        self.generic_visit(node) 

def test_importfrom(self):
        imp = ast.ImportFrom(None, [ast.alias("x", None)], -42)
        self.stmt(imp, "level less than -1")
        self.stmt(ast.ImportFrom(None, [], 0), "empty names on ImportFrom") 

def sort_imports(node):
    """sorts the node in place"""
    imports = [stmnt for stmnt in node.body 
              if isinstance(stmnt, Import)]
    imports = [name for imp in imports for name in imp.names]
    import_froms = [stmnt for stmnt in node.body 
                   if isinstance(stmnt, ImportFrom)]
    remainder = [stmnt for stmnt in node.body 
                if not (stmnt in imports or stmnt in import_froms)]
    imports.sort(key=lambda i: i.name)
    import_froms.sort(key=lambda i: i.module)
    for impfrm in import_froms:
        impfrm.names.sort(key=lambda i: i.name)
    # print([alias.name for alias in import_froms[-1].names])
    node.body = import_froms + imports + remainder 

def test_non_interned_future_from_ast(self):
        mod = ast.parse("from __future__ import division")
        self.assertIsInstance(mod.body[0], ast.ImportFrom)
        mod.body[0].module = " __future__ ".strip()
        compile(mod, "<test>", "exec") 

def list_import_names(tree, package_name):
    for node in ast.iter_child_nodes(tree):
        if isinstance(node, ast.Import):
            for name in node.names:
                yield ImportItem(
                    original_name=name.name,
                    resolved_name=name.name,
                    lineno=node.lineno,
                    col_offset=node.col_offset,
                    relative=False
                )
        elif isinstance(node, ast.ImportFrom):
            for name in node.names:
                original_name = (
                    '.' * node.level +
                    (node.module + '.' if node.module else '') +
                    name.name
                )
                yield ImportItem(
                    original_name=original_name,
                    resolved_name=resolve_name(original_name, package_name),
                    lineno=node.lineno,
                    col_offset=node.col_offset,
                    relative=True
                )
        else:
            continue 

def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        for name in node.names:
            self.read(name.name, node)
            if name.asname:
                self.define(name.asname, node)

        self.generic_visit(node) 

def test_bad_integer(self):
        # issue13436: Bad error message with invalid numeric values
        body = [ast.ImportFrom(module='time',
                               names=[ast.alias(name='sleep')],
                               level=None,
                               lineno=None, col_offset=None)]
        mod = ast.Module(body)
        with self.assertRaises(ValueError) as cm:
            compile(mod, 'test', 'exec')
        self.assertIn("invalid integer value: None", str(cm.exception)) 

def add_closure_inspection(f):
    @wraps(f)
    def wrapper(pickler: '_EboniteRequirementAnalyzer', obj):
        closure = inspect.getclosurevars(obj)
        for field in ['nonlocals', 'globals']:
            for o in getattr(closure, field).values():
                if isinstance(o, ModuleType):
                    pickler._add_requirement(o)
                else:
                    pickler.save(o)

        if is_from_installable_module(obj):
            return f(pickler, obj)

        # to add from local imports inside user (non PIP package) code
        tree = ast.parse(inspect.getsource(obj).strip())

        class ImportFromVisitor(ast.NodeVisitor):
            def visit_ImportFrom(self, node: ast.ImportFrom):  # noqa
                warnings.warn(f'Detected local import in {obj.__module__}.{obj.__name__}')
                if node.level == 0:
                    mod = import_module(node.module)
                else:
                    mod = import_module('.' + node.module, get_object_module(obj).__package__)
                pickler._add_requirement(mod)

        ImportFromVisitor().visit(tree)

        return f(pickler, obj)

    return wrapper 

def get_all_import_names_mentioned_in_import(tree):
    import_names = []
    imports = get_all_imports(tree)
    for import_node in imports:
        if isinstance(import_node, ast.ImportFrom):
            import_names.append(import_node.module)
        elif isinstance(import_node, ast.Import):
            import_names += [import_object.name for import_object in import_node.names]
    return import_names 

def get_all_imports(root):
    return get_nodes_of_type(root, (ast.ImportFrom, ast.Import)) 

def process_script(self, path):
        # side effect: updates set self.externals and dict self.script_imports (adding value only for key = path)
        # returns the script_imports of the processed script
        with open(path, 'r') as f:
            code = f.read()

        tree = ast.parse(code)

        self.script_imports[path] = set([])
        for node in tree.body:
            if node.__class__ is ast.Import:
                module_names = [alias.name for alias in node.names]

            elif node.__class__ is ast.ImportFrom:
                parent_module_name = node.module

                module_names = ['{}.{}'.format(parent_module_name, alias.name) for alias in node.names]

            else:
                continue

            for module_name in module_names:
                if module_name.startswith(self.package_name):
                    self.script_imports[path].add(module_name)
                else:
                    self.externals.add(module_name)

        return self.script_imports[path] 

def parse(code, nas_mode=None):
    """Annotate user code.
    Return annotated code (str) if annotation detected; return None if not.
    code: original user code (str),
    nas_mode: the mode of NAS given that NAS interface is used
    """
    try:
        ast_tree = ast.parse(code)
    except Exception:
        raise RuntimeError('Bad Python code')

    transformer = Transformer(nas_mode)
    try:
        transformer.visit(ast_tree)
    except AssertionError as exc:
        raise RuntimeError('%d: %s' % (ast_tree.last_line, exc.args[0]))

    if not transformer.annotated:
        return None

    last_future_import = -1
    import_nni = ast.Import(names=[ast.alias(name='nni', asname=None)])
    nodes = ast_tree.body
    for i, _ in enumerate(nodes):
        if type(nodes[i]) is ast.ImportFrom and nodes[i].module == '__future__':
            last_future_import = i
    nodes.insert(last_future_import + 1, import_nni)
    # enas, oneshot and darts modes for tensorflow need tensorflow module, so we import it here
    if nas_mode in ['enas_mode', 'oneshot_mode', 'darts_mode']:
        import_tf = ast.Import(names=[ast.alias(name='tensorflow', asname=None)])
        nodes.insert(last_future_import + 1, import_tf)

    return astor.to_source(ast_tree) 

def is_future_import(node):
    return isinstance(node, ast.ImportFrom) and node.module == "__future__" 

def iter_imports(path):
    '''Yield imports in *path*'''

    for node in ast.parse(open(path, 'rb').read()).body:
        if isinstance(node, ast.ImportFrom):
            if node.module is None:
                prefix = ()
            else:
                prefix = tuple(node.module.split('.'))
            for snode in node.names:
                yield (node.level, prefix + (snode.name,))
        elif isinstance(node, ast.Import):
            for node in node.names:
                yield (0, tuple(node.name.split('.'))) 

def get_all_imported_names_from_tree(tree):
    imported_names = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.ImportFrom):
            continue
        for name in node.names:
            imported_name = name.asname or name.name
            imported_names.append(imported_name)
    return set(imported_names)