提交 8841a9c3 编写于 作者: C Cleber Rosa

safeloader: remove duplicate code by introducing statement_import_as()

And reusing them in two different locations.
Signed-off-by: NCleber Rosa <crosa@redhat.com>
上级 42875fa7
......@@ -62,18 +62,33 @@ class AvocadoModule(object):
# Looking for a 'import avocado'
elif isinstance(statement, ast.Import):
for name in statement.names:
if name.name == 'avocado':
if name.asname is not None:
self.mod_imports.add(name.asname)
else:
self.mod_imports.add(name.name)
imp_name = statement_import_as(statement).get('avocado', None)
if imp_name is not None:
self.mod_imports.add(imp_name)
# Looking for a 'class Anything(anything):'
elif isinstance(statement, ast.ClassDef):
yield statement
def statement_import_as(statement):
"""
Returns a mapping of imported module names whether using aliases or not
:param statement: an AST import statement
:type statement: ast.Import
:returns: a mapping of names {<realname>: <alias>} of modules imported
:rtype: dict
"""
result = {}
for name in statement.names:
if name.asname is not None:
result[name.name] = name.asname
else:
result[name.name] = name.name
return result
def modules_imported_as(module):
"""
Returns a mapping of imported module names whether using aliases or not
......@@ -105,13 +120,8 @@ def modules_imported_as(module):
"""
result = {}
for statement in module.body:
# Looking for a 'import <module>'
if isinstance(statement, ast.Import):
for name in statement.names:
if name.asname is not None:
result[name.name] = name.asname
else:
result[name.name] = name.name
result.update(statement_import_as(statement))
return result
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册