"""Astroid hooks for various builtins."""
import sys
from functools import partial
from textwrap import dedent
import six
from astroid import (MANAGER, UseInferenceDefault,
inference_tip, YES, InferenceError, UnresolvableName)
from astroid import nodes
from astroid.builder import AstroidBuilder
def _extend_str(class_node, rvalue):
"""function to extend builtin str/unicode class"""
# TODO(cpopa): this approach will make astroid to believe
# that some arguments can be passed by keyword, but
# unfortunately, strings and bytes don't accept keyword arguments.
code = dedent('''
class whatever(object):
def join(self, iterable):
return {rvalue}
def replace(self, old, new, count=None):
return {rvalue}
def format(self, *args, **kwargs):
return {rvalue}
def encode(self, encoding='ascii', errors=None):
return ''
def decode(self, encoding='ascii', errors=None):
return u''
def capitalize(self):
return {rvalue}
def title(self):
return {rvalue}
def lower(self):
return {rvalue}
def upper(self):
return {rvalue}
def swapcase(self):
return {rvalue}
def index(self, sub, start=None, end=None):
return 0
def find(self, sub, start=None, end=None):
return 0
def count(self, sub, start=None, end=None):
return 0
def strip(self, chars=None):
return {rvalue}
def lstrip(self, chars=None):
return {rvalue}
def rstrip(self, chars=None):
return {rvalue}
def rjust(self, width, fillchar=None):
return {rvalue}
def center(self, width, fillchar=None):
return {rvalue}
def ljust(self, width, fillchar=None):
return {rvalue}
''')
code = code.format(rvalue=rvalue)
fake = AstroidBuilder(MANAGER).string_build(code)['whatever']
for method in fake.mymethods():
class_node.locals[method.name] = [method]
method.parent = class_node
def extend_builtins(class_transforms):
from astroid.bases import BUILTINS
builtin_ast = MANAGER.astroid_cache[BUILTINS]
for class_name, transform in class_transforms.items():
transform(builtin_ast[class_name])
if sys.version_info > (3, 0):
extend_builtins({'bytes': partial(_extend_str, rvalue="b''"),
'str': partial(_extend_str, rvalue="''")})
else:
extend_builtins({'str': partial(_extend_str, rvalue="''"),
'unicode': partial(_extend_str, rvalue="u''")})
def register_builtin_transform(transform, builtin_name):
"""Register a new transform function for the given *builtin_name*.
The transform function must accept two parameters, a node and
an optional context.
"""
def _transform_wrapper(node, context=None):
result = transform(node, context=context)
if result:
result.parent = node
result.lineno = node.lineno
result.col_offset = node.col_offset
return iter([result])
MANAGER.register_transform(nodes.CallFunc,
inference_tip(_transform_wrapper),
lambda n: (isinstance(n.func, nodes.Name) and
n.func.name == builtin_name))
def _generic_inference(node, context, node_type, transform):
args = node.args
if not args:
return node_type()
if len(node.args) > 1:
raise UseInferenceDefault()
arg, = args
transformed = transform(arg)
if not transformed:
try:
infered = next(arg.infer(context=context))
except (InferenceError, StopIteration):
raise UseInferenceDefault()
if infered is YES:
raise UseInferenceDefault()
transformed = transform(infered)
if not transformed or transformed is YES:
raise UseInferenceDefault()
return transformed
def _generic_transform(arg, klass, iterables, build_elts):
if isinstance(arg, klass):
return arg
elif isinstance(arg, iterables):
if not all(isinstance(elt, nodes.Const)
for elt in arg.elts):
# TODO(cpopa): Don't support heterogenous elements.
# Not yet, though.
raise UseInferenceDefault()
elts = [elt.value for elt in arg.elts]
elif isinstance(arg, nodes.Dict):
if not all(isinstance(elt[0], nodes.Const)
for elt in arg.items):
raise UseInferenceDefault()
elts = [item[0].value for item in arg.items]
elif (isinstance(arg, nodes.Const) and
isinstance(arg.value, (six.string_types, six.binary_type))):
elts = arg.value
else:
return
return klass(elts=build_elts(elts))
def _infer_builtin(node, context,
klass=None, iterables=None,
build_elts=None):
transform_func = partial(
_generic_transform,
klass=klass,
iterables=iterables,
build_elts=build_elts)
return _generic_inference(node, context, klass, transform_func)
# pylint: disable=invalid-name
infer_tuple = partial(
_infer_builtin,
klass=nodes.Tuple,
iterables=(nodes.List, nodes.Set),
build_elts=tuple)
infer_list = partial(
_infer_builtin,
klass=nodes.List,
iterables=(nodes.Tuple, nodes.Set),
build_elts=list)
infer_set = partial(
_infer_builtin,
klass=nodes.Set,
iterables=(nodes.List, nodes.Tuple),
build_elts=set)
def _get_elts(arg, context):
is_iterable = lambda n: isinstance(n,
(nodes.List, nodes.Tuple, nodes.Set))
try:
infered = next(arg.infer(context))
except (InferenceError, UnresolvableName):
raise UseInferenceDefault()
if isinstance(infered, nodes.Dict):
items = infered.items
elif is_iterable(infered):
items = []
for elt in infered.elts:
# If an item is not a pair of two items,
# then fallback to the default inference.
# Also, take in consideration only hashable items,
# tuples and consts. We are choosing Names as well.
if not is_iterable(elt):
raise UseInferenceDefault()
if len(elt.elts) != 2:
raise UseInferenceDefault()
if not isinstance(elt.elts[0],
(nodes.Tuple, nodes.Const, nodes.Name)):
raise UseInferenceDefault()
items.append(tuple(elt.elts))
else:
raise UseInferenceDefault()
return items
def infer_dict(node, context=None):
"""Try to infer a dict call to a Dict node.
The function treats the following cases:
* dict()
* dict(mapping)
* dict(iterable)
* dict(iterable, **kwargs)
* dict(mapping, **kwargs)
* dict(**kwargs)
If a case can't be infered, we'll fallback to default inference.
"""
has_keywords = lambda args: all(isinstance(arg, nodes.Keyword)
for arg in args)
if not node.args and not node.kwargs:
# dict()
return nodes.Dict()
elif has_keywords(node.args) and node.args:
# dict(a=1, b=2, c=4)
items = [(nodes.Const(arg.arg), arg.value) for arg in node.args]
elif (len(node.args) >= 2 and
has_keywords(node.args[1:])):
# dict(some_iterable, b=2, c=4)
elts = _get_elts(node.args[0], context)
keys = [(nodes.Const(arg.arg), arg.value) for arg in node.args[1:]]
items = elts + keys
elif len(node.args) == 1:
items = _get_elts(node.args[0], context)
else:
raise UseInferenceDefault()
empty = nodes.Dict()
empty.items = items
return empty
# Builtins inference
register_builtin_transform(infer_tuple, 'tuple')
register_builtin_transform(infer_set, 'set')
register_builtin_transform(infer_list, 'list')
register_builtin_transform(infer_dict, 'dict')
Codebase Browser
chromium