Merge with the linter branch (especially the changes of pep484.

jedi/_compatibility.py

@@ -13,6 +13,7 @@ except ImportError:
 
 is_py3 = sys.version_info[0] >= 3
 is_py33 = is_py3 and sys.version_info.minor >= 3
+is_py35 = is_py3 and sys.version_info.minor >= 5
 is_py26 = not is_py3 and sys.version_info[1] < 7
 
 

jedi/api/keywords.py

@@ -1,7 +1,7 @@
 import pydoc
 import keyword
 
-from jedi._compatibility import is_py3
+from jedi._compatibility import is_py3, is_py35
 from jedi import common
 from jedi.evaluate.helpers import FakeName
 from jedi.parser.tree import Leaf
@@ -12,7 +12,12 @@ except ImportError:
     import pydoc_topics
 
 if is_py3:
-    keys = keyword.kwlist
+    if is_py35:
+        # in python 3.5 async and await are not proper keywords, but for
+        # completion pursposes should as as though they are
+        keys = keyword.kwlist + ["async", "await"]
+    else:
+        keys = keyword.kwlist
 else:
     keys = keyword.kwlist + ['None', 'False', 'True']
 

jedi/debug.py

@@ -35,7 +35,12 @@ try:
                 # after every output the stream is reset automatically we don't
                 # need this.
                 initialise.atexit_done = True
-                init()
+                try:
+                    init()
+                except Exception:
+                    # Colorama fails with initializing under vim and is buggy in
+                    # version 0.3.6.
+                    pass
             _inited = True
 
 except ImportError:

jedi/evaluate/__init__.py

@@ -292,7 +292,7 @@ class Evaluator(object):
             types = set([element])  # TODO this is no real evaluation.
         elif element.type == 'expr_stmt':
             types = self.eval_statement(element)
-        elif element.type == 'power':
+        elif element.type in ('power', 'atom_expr'):
             types = self._eval_atom(element.children[0])
             for trailer in element.children[1:]:
                 if trailer == '**':  # has a power operation.
@@ -392,8 +392,7 @@ class Evaluator(object):
 
         new_types = set()
         if trailer_op == '[':
-            for trailer_typ in iterable.create_index_types(self, node):
-                new_types |= iterable.py__getitem__(self, types, trailer_typ, trailer_op)
+            new_types |= iterable.py__getitem__(self, types, trailer)
         else:
             for typ in types:
                 debug.dbg('eval_trailer: %s in scope %s', trailer, typ)

jedi/evaluate/analysis.py

@@ -174,7 +174,7 @@ def _check_for_exception_catch(evaluator, jedi_obj, exception, payload=None):
     def check_hasattr(node, suite):
         try:
             assert suite.start_pos <= jedi_obj.start_pos < suite.end_pos
-            assert node.type == 'power'
+            assert node.type in ('power', 'atom_expr')
             base = node.children[0]
             assert base.type == 'name' and base.value == 'hasattr'
             trailer = node.children[1]

jedi/evaluate/dynamic.py

@@ -89,7 +89,7 @@ def search_function_call(evaluator, func):
                     parent = parent.parent
 
                 trailer = None
-                if tree.is_node(parent, 'power'):
+                if tree.is_node(parent, 'power', 'atom_expr'):
                     for t in parent.children[1:]:
                         if t == '**':
                             break

jedi/evaluate/finder.py

@@ -311,6 +311,14 @@ class NameFinder(object):
 @memoize_default(set(), evaluator_is_first_arg=True)
 def _name_to_types(evaluator, name, scope):
     typ = name.get_definition()
+    if typ.isinstance(tree.ForStmt):
+        types = pep0484.find_type_from_comment_hint_for(evaluator, typ, name)
+        if types:
+            return types
+    if typ.isinstance(tree.WithStmt):
+        types = pep0484.find_type_from_comment_hint_with(evaluator, typ, name)
+        if types:
+            return types
     if typ.isinstance(tree.ForStmt, tree.CompFor):
         container_types = evaluator.eval_element(typ.children[3])
         for_types = iterable.py__iter__types(evaluator, container_types, typ.children[3])
@@ -362,6 +370,10 @@ def _remove_statements(evaluator, stmt, name):
         check_instance = stmt.instance
         stmt = stmt.var
 
+    pep0484types = \
+        pep0484.find_type_from_comment_hint_assign(evaluator, stmt, name)
+    if pep0484types:
+        return pep0484types
     types |= evaluator.eval_statement(stmt, seek_name=name)
 
     if check_instance is not None:
@@ -455,7 +467,7 @@ def check_flow_information(evaluator, flow, search_name, pos):
 
 def _check_isinstance_type(evaluator, element, search_name):
     try:
-        assert element.type == 'power'
+        assert element.type in ('power', 'atom_expr')
         # this might be removed if we analyze and, etc
         assert len(element.children) == 2
         first, trailer = element.children

jedi/evaluate/helpers.py

@@ -89,7 +89,8 @@ def call_of_name(name, cut_own_trailer=False):
             return name
 
     power = par.parent
-    if tree.is_node(power, 'power') and power.children[0] != name \
+    if tree.is_node(power, 'power', 'atom_expr') \
+            and power.children[0] != name \
             and not (power.children[-2] == '**' and
                      name.start_pos > power.children[-1].start_pos):
         par = power

jedi/evaluate/iterable.py

@@ -29,6 +29,7 @@ from jedi.evaluate import compiled
 from jedi.evaluate import helpers
 from jedi.evaluate.cache import CachedMetaClass, memoize_default
 from jedi.evaluate import analysis
+from jedi.evaluate import pep0484
 
 
 class IterableWrapper(tree.Base):
@@ -430,6 +431,9 @@ class ImplicitTuple(_FakeArray):
 
 class FakeSequence(_FakeArray):
     def __init__(self, evaluator, sequence_values, type):
+        """
+        type should be one of "tuple", "list"
+        """
         super(FakeSequence, self).__init__(evaluator, sequence_values, type)
         self._sequence_values = sequence_values
 
@@ -524,7 +528,7 @@ def unpack_tuple_to_dict(evaluator, types, exprlist):
             analysis.add(evaluator, 'value-error-too-few-values', has_parts,
                          message="ValueError: need more than %s values to unpack" % n)
         return dct
-    elif exprlist.type == 'power':
+    elif exprlist.type == 'power' or exprlist.type == 'atom_expr':
         # Something like ``arr[x], var = ...``.
         # This is something that is not yet supported, would also be difficult
         # to write into a dict.
@@ -559,37 +563,56 @@ def py__iter__types(evaluator, types, node=None):
     return unite(py__iter__(evaluator, types, node))
 
 
-def py__getitem__(evaluator, types, index, node):
+def py__getitem__(evaluator, types, trailer):
+    from jedi.evaluate.representation import Class
     result = set()
 
-    # Index handling.
-    if isinstance(index, (compiled.CompiledObject, Slice)):
-        index = index.obj
+    trailer_op, node, trailer_cl = trailer.children
+    assert trailer_op == "["
+    assert trailer_cl == "]"
 
-    if type(index) not in (float, int, str, unicode, slice):
-        # If the index is not clearly defined, we have to get all the
-        # possiblities.
-        for typ in list(types):
-            if isinstance(typ, Array) and typ.type == 'dict':
+    # special case: PEP0484 typing module, see
+    # https://github.com/davidhalter/jedi/issues/663
+    for typ in list(types):
+        if isinstance(typ, Class):
+            typing_module_types = \
+                pep0484.get_types_for_typing_module(evaluator, typ, node)
+            if typing_module_types is not None:
                 types.remove(typ)
-                result |= typ.dict_values()
-        return result | py__iter__types(evaluator, types)
+                result |= typing_module_types
 
-    for typ in types:
-        # The actual getitem call.
-        try:
-            getitem = typ.py__getitem__
-        except AttributeError:
-            analysis.add(evaluator, 'type-error-not-subscriptable', node,
-                         message="TypeError: '%s' object is not subscriptable" % typ)
-        else:
+    if not types:
+        # all consumed by special cases
+        return result
+
+    for index in create_index_types(evaluator, node):
+        if isinstance(index, (compiled.CompiledObject, Slice)):
+            index = index.obj
+
+        if type(index) not in (float, int, str, unicode, slice):
+            # If the index is not clearly defined, we have to get all the
+            # possiblities.
+            for typ in list(types):
+                if isinstance(typ, Array) and typ.type == 'dict':
+                    types.remove(typ)
+                    result |= typ.dict_values()
+            return result | py__iter__types(evaluator, types)
+
+        for typ in types:
+            # The actual getitem call.
             try:
-                result |= getitem(index)
-            except IndexError:
-                result |= py__iter__types(evaluator, set([typ]))
-            except KeyError:
-                # Must be a dict. Lists don't raise IndexErrors.
-                result |= typ.dict_values()
+                getitem = typ.py__getitem__
+            except AttributeError:
+                analysis.add(evaluator, 'type-error-not-subscriptable', trailer_op,
+                             message="TypeError: '%s' object is not subscriptable" % typ)
+            else:
+                try:
+                    result |= getitem(index)
+                except IndexError:
+                    result |= py__iter__types(evaluator, set([typ]))
+                except KeyError:
+                    # Must be a dict. Lists don't raise KeyErrors.
+                    result |= typ.dict_values()
     return result
 
 

jedi/evaluate/jedi_typing.py

@@ -0,0 +1,100 @@
+"""
+This module is not intended to be used in jedi, rather it will be fed to the
+jedi-parser to replace classes in the typing module
+"""
+
+try:
+    from collections import abc
+except ImportError:
+    # python 2
+    import collections as abc
+
+
+def factory(typing_name, indextypes):
+    class Iterable(abc.Iterable):
+        def __iter__(self):
+            while True:
+                yield indextypes[0]()
+
+    class Iterator(Iterable, abc.Iterator):
+        def next(self):
+            """ needed for python 2 """
+            return self.__next__()
+
+        def __next__(self):
+            return indextypes[0]()
+
+    class Sequence(abc.Sequence):
+        def __getitem__(self, index):
+            return indextypes[0]()
+
+    class MutableSequence(Sequence, abc.MutableSequence):
+        pass
+
+    class List(MutableSequence, list):
+        pass
+
+    class Tuple(Sequence, tuple):
+        def __getitem__(self, index):
+            if indextypes[1] == Ellipsis:
+                # https://www.python.org/dev/peps/pep-0484/#the-typing-module
+                # Tuple[int, ...] means a tuple of ints of indetermined length
+                return indextypes[0]()
+            else:
+                return indextypes[index]()
+
+    class AbstractSet(Iterable, abc.Set):
+        pass
+
+    class MutableSet(AbstractSet, abc.MutableSet):
+        pass
+
+    class KeysView(Iterable, abc.KeysView):
+        pass
+
+    class ValuesView(abc.ValuesView):
+        def __iter__(self):
+            while True:
+                yield indextypes[1]()
+
+    class ItemsView(abc.ItemsView):
+        def __iter__(self):
+            while True:
+                yield (indextypes[0](), indextypes[1]())
+
+    class Mapping(Iterable, abc.Mapping):
+        def __getitem__(self, item):
+            return indextypes[1]()
+
+        def keys(self):
+            return KeysView()
+
+        def values(self):
+            return ValuesView()
+
+        def items(self):
+            return ItemsView()
+
+    class MutableMapping(Mapping, abc.MutableMapping):
+        pass
+
+    class Dict(MutableMapping, dict):
+        pass
+
+    dct = {
+        "Sequence": Sequence,
+        "MutableSequence": MutableSequence,
+        "List": List,
+        "Iterable": Iterable,
+        "Iterator": Iterator,
+        "AbstractSet": AbstractSet,
+        "MutableSet": MutableSet,
+        "Mapping": Mapping,
+        "MutableMapping": MutableMapping,
+        "Tuple": Tuple,
+        "KeysView": KeysView,
+        "ItemsView": ItemsView,
+        "ValuesView": ValuesView,
+        "Dict": Dict,
+    }
+    return dct[typing_name]

jedi/evaluate/param.py

@@ -47,7 +47,10 @@ class Arguments(tree.Base):
             for el in self.argument_node:
                 yield 0, el
         else:
-            if not tree.is_node(self.argument_node, 'arglist'):
+            if not (tree.is_node(self.argument_node, 'arglist') or (
+                    # in python 3.5 **arg is an argument, not arglist
+                    (tree.is_node(self.argument_node, 'argument') and
+                     self.argument_node.children[0] in ('*', '**')))):
                 yield 0, self.argument_node
                 return
 
@@ -57,6 +60,10 @@ class Arguments(tree.Base):
                     continue
                 elif child in ('*', '**'):
                     yield len(child.value), next(iterator)
+                elif tree.is_node(child, 'argument') and \
+                        child.children[0] in ('*', '**'):
+                    assert len(child.children) == 2
+                    yield len(child.children[0].value), child.children[1]
                 else:
                     yield 0, child
 

jedi/evaluate/pep0484.py

@@ -9,47 +9,74 @@ v Function parameter annotations with builtin/custom type classes
 v Function returntype annotations with builtin/custom type classes
 v Function parameter annotations with strings (forward reference)
 v Function return type annotations with strings (forward reference)
-x Local variable type hints
+v Local variable type hints
 v Assigned types: `Url = str\ndef get(url:Url) -> str:`
-x Type hints in `with` statements
+v Type hints in `with` statements
 x Stub files support
 x support `@no_type_check` and `@no_type_check_decorator`
-x support for type hint comments `# type: (int, str) -> int`. See comment from
-    Guido https://github.com/davidhalter/jedi/issues/662
+x support for typing.cast() operator
+x support for type hint comments for functions, `# type: (int, str) -> int`.
+    See comment from Guido https://github.com/davidhalter/jedi/issues/662
 """
 
-from itertools import chain
+import itertools
 
-from jedi.parser import Parser, load_grammar, ParseError
+import os
+from jedi.parser import \
+    Parser, load_grammar, ParseError, ParserWithRecovery, tree
 from jedi.evaluate.cache import memoize_default
-from jedi.evaluate.compiled import CompiledObject
+from jedi.common import unite
+from jedi.evaluate import compiled
 from jedi import debug
+from jedi import _compatibility
+import re
 
 
-def _evaluate_for_annotation(evaluator, annotation):
+def _evaluate_for_annotation(evaluator, annotation, index=None):
+    """
+    Evaluates a string-node, looking for an annotation
+    If index is not None, the annotation is expected to be a tuple
+    and we're interested in that index
+    """
     if annotation is not None:
-        definitions = set()
-        for definition in evaluator.eval_element(annotation):
-            if (isinstance(definition, CompiledObject) and
-                    isinstance(definition.obj, str)):
-                try:
-                    p = Parser(load_grammar(), definition.obj, start_symbol='eval_input')
-                    element = p.get_parsed_node()
-                except ParseError:
-                    debug.warning('Annotation not parsed: %s' % definition.obj)
-                else:
-                    module = annotation.get_parent_until()
-                    p.position_modifier.line = module.end_pos[0]
-                    element.parent = module
-                    definitions |= evaluator.eval_element(element)
-            else:
-                definitions.add(definition)
-        return list(chain.from_iterable(
+        definitions = evaluator.eval_element(
+            _fix_forward_reference(evaluator, annotation))
+        if index is not None:
+            definitions = list(itertools.chain.from_iterable(
+                definition.py__getitem__(index) for definition in definitions
+                if definition.type == 'tuple' and
+                len(list(definition.py__iter__())) >= index))
+        return list(itertools.chain.from_iterable(
             evaluator.execute(d) for d in definitions))
     else:
         return []
 
 
+def _fix_forward_reference(evaluator, node):
+    evaled_nodes = evaluator.eval_element(node)
+    if len(evaled_nodes) != 1:
+        debug.warning("Eval'ed typing index %s should lead to 1 object, "
+                      " not %s" % (node, evaled_nodes))
+        return node
+    evaled_node = list(evaled_nodes)[0]
+    if isinstance(evaled_node, compiled.CompiledObject) and \
+            isinstance(evaled_node.obj, str):
+        try:
+            p = Parser(load_grammar(), _compatibility.unicode(evaled_node.obj),
+                       start_symbol='eval_input')
+            newnode = p.get_parsed_node()
+        except ParseError:
+            debug.warning('Annotation not parsed: %s' % evaled_node.obj)
+            return node
+        else:
+            module = node.get_parent_until()
+            p.position_modifier.line = module.end_pos[0]
+            newnode.parent = module
+            return newnode
+    else:
+        return node
+
+
 @memoize_default(None, evaluator_is_first_arg=True)
 def follow_param(evaluator, param):
     annotation = param.annotation()
@@ -60,3 +87,109 @@ def follow_param(evaluator, param):
 def find_return_types(evaluator, func):
     annotation = func.py__annotations__().get("return", None)
     return _evaluate_for_annotation(evaluator, annotation)
+
+
+_typing_module = None
+
+
+def _get_typing_replacement_module():
+    """
+    The idea is to return our jedi replacement for the PEP-0484 typing module
+    as discussed at https://github.com/davidhalter/jedi/issues/663
+    """
+    global _typing_module
+    if _typing_module is None:
+        typing_path = \
+            os.path.abspath(os.path.join(__file__, "../jedi_typing.py"))
+        with open(typing_path) as f:
+            code = _compatibility.unicode(f.read())
+        p = ParserWithRecovery(load_grammar(), code)
+        _typing_module = p.module
+    return _typing_module
+
+
+def get_types_for_typing_module(evaluator, typ, node):
+    from jedi.evaluate.iterable import FakeSequence
+    if not typ.base.get_parent_until().name.value == "typing":
+        return None
+    # we assume that any class using [] in a module called
+    # "typing" with a name for which we have a replacement
+    # should be replaced by that class. This is not 100%
+    # airtight but I don't have a better idea to check that it's
+    # actually the PEP-0484 typing module and not some other
+    if tree.is_node(node, "subscriptlist"):
+        nodes = node.children[::2]  # skip the commas
+    else:
+        nodes = [node]
+    del node
+
+    nodes = [_fix_forward_reference(evaluator, node) for node in nodes]
+
+    # hacked in Union and Optional, since it's hard to do nicely in parsed code
+    if typ.name.value == "Union":
+        return unite(evaluator.eval_element(node) for node in nodes)
+    if typ.name.value == "Optional":
+        return evaluator.eval_element(nodes[0])
+
+    typing = _get_typing_replacement_module()
+    factories = evaluator.find_types(typing, "factory")
+    assert len(factories) == 1
+    factory = list(factories)[0]
+    assert factory
+    function_body_nodes = factory.children[4].children
+    valid_classnames = set(child.name.value
+                           for child in function_body_nodes
+                           if isinstance(child, tree.Class))
+    if typ.name.value not in valid_classnames:
+        return None
+    compiled_classname = compiled.create(evaluator, typ.name.value)
+
+    args = FakeSequence(evaluator, nodes, "tuple")
+
+    result = evaluator.execute_evaluated(factory, compiled_classname, args)
+    return result
+
+
+def find_type_from_comment_hint_for(evaluator, node, name):
+    return \
+        _find_type_from_comment_hint(evaluator, node, node.children[1], name)
+
+
+def find_type_from_comment_hint_with(evaluator, node, name):
+    assert len(node.children[1].children) == 3, \
+        "Can only be here when children[1] is 'foo() as f'"
+    return _find_type_from_comment_hint(
+        evaluator, node, node.children[1].children[2], name)
+
+
+def find_type_from_comment_hint_assign(evaluator, node, name):
+    return \
+        _find_type_from_comment_hint(evaluator, node, node.children[0], name)
+
+
+def _find_type_from_comment_hint(evaluator, node, varlist, name):
+    index = None
+    if varlist.type in ("testlist_star_expr", "exprlist"):
+        # something like "a, b = 1, 2"
+        index = 0
+        for child in varlist.children:
+            if child == name:
+                break
+            if child.type == "operator":
+                continue
+            index += 1
+        else:
+            return []
+
+    comment = node.get_following_comment_same_line()
+    if comment is None:
+        return []
+    match = re.match(r"^#\s*type:\s*([^#]*)", comment)
+    if not match:
+        return []
+    annotation = tree.String(
+        tree.zero_position_modifier,
+        repr(str(match.group(1).strip())),
+        node.start_pos)
+    annotation.parent = node.parent
+    return _evaluate_for_annotation(evaluator, annotation, index)

jedi/evaluate/sys_path.py

@@ -99,7 +99,8 @@ def _paths_from_assignment(evaluator, expr_stmt):
     for assignee, operator in zip(expr_stmt.children[::2], expr_stmt.children[1::2]):
         try:
             assert operator in ['=', '+=']
-            assert tree.is_node(assignee, 'power') and len(assignee.children) > 1
+            assert tree.is_node(assignee, 'power', 'atom_expr') and \
+                len(assignee.children) > 1
             c = assignee.children
             assert c[0].type == 'name' and c[0].value == 'sys'
             trailer = c[1]
@@ -152,7 +153,7 @@ def _check_module(evaluator, module):
     def get_sys_path_powers(names):
         for name in names:
             power = name.parent.parent
-            if tree.is_node(power, 'power'):
+            if tree.is_node(power, 'power', 'atom_expr'):
                 c = power.children
                 if isinstance(c[0], tree.Name) and c[0].value == 'sys' \
                         and tree.is_node(c[1], 'trailer'):

jedi/parser/__init__.py

@@ -46,7 +46,7 @@ def load_grammar(version='3.4'):
     if version in ('3.2', '3.3'):
         version = '3.4'
     elif version == '2.6':
-        version == '2.7'
+        version = '2.7'
 
     file = 'grammar' + version + '.txt'
 

jedi/parser/grammar3.5.txt

@@ -15,15 +15,17 @@
 #       file_input is a module or sequence of commands read from an input file;
 #       eval_input is the input for the eval() functions.
 # NB: compound_stmt in single_input is followed by extra NEWLINE!
-single_input: NEWLINE | simple_stmt | compound_stmt NEWLINE
 file_input: (NEWLINE | stmt)* ENDMARKER
+single_input: NEWLINE | simple_stmt | compound_stmt NEWLINE
 eval_input: testlist NEWLINE* ENDMARKER
 
 decorator: '@' dotted_name [ '(' [arglist] ')' ] NEWLINE
 decorators: decorator+
 decorated: decorators (classdef | funcdef | async_funcdef)
 
-async_funcdef: ASYNC funcdef
+# NOTE: Reinoud Elhorst, using ASYNC/AWAIT keywords instead of tokens
+# skipping python3.5 compatibility, in favour of 3.7 solution
+async_funcdef: 'async' funcdef
 funcdef: 'def' NAME parameters ['->' test] ':' suite
 
 parameters: '(' [typedargslist] ')'
@@ -69,7 +71,7 @@ nonlocal_stmt: 'nonlocal' NAME (',' NAME)*
 assert_stmt: 'assert' test [',' test]
 
 compound_stmt: if_stmt | while_stmt | for_stmt | try_stmt | with_stmt | funcdef | classdef | decorated | async_stmt
-async_stmt: ASYNC (funcdef | with_stmt | for_stmt)
+async_stmt: 'async' (funcdef | with_stmt | for_stmt)
 if_stmt: 'if' test ':' suite ('elif' test ':' suite)* ['else' ':' suite]
 while_stmt: 'while' test ':' suite ['else' ':' suite]
 for_stmt: 'for' exprlist 'in' testlist ':' suite ['else' ':' suite]
@@ -82,7 +84,9 @@ with_stmt: 'with' with_item (',' with_item)*  ':' suite
 with_item: test ['as' expr]
 # NB compile.c makes sure that the default except clause is last
 except_clause: 'except' [test ['as' NAME]]
-suite: simple_stmt | NEWLINE INDENT stmt+ DEDENT
+# Edit by David Halter: The stmt is now optional. This reflects how Jedi allows
+# classes and functions to be empty, which is beneficial for autocompletion.
+suite: simple_stmt | NEWLINE INDENT stmt* DEDENT
 
 test: or_test ['if' or_test 'else' test] | lambdef
 test_nocond: or_test | lambdef_nocond
@@ -104,7 +108,7 @@ arith_expr: term (('+'|'-') term)*
 term: factor (('*'|'@'|'/'|'%'|'//') factor)*
 factor: ('+'|'-'|'~') factor | power
 power: atom_expr ['**' factor]
-atom_expr: [AWAIT] atom trailer*
+atom_expr: ['await'] atom trailer*
 atom: ('(' [yield_expr|testlist_comp] ')' |
        '[' [testlist_comp] ']' |
        '{' [dictorsetmaker] '}' |

jedi/parser/token.py

@@ -1,6 +1,6 @@
 from __future__ import absolute_import
 
-from jedi._compatibility import is_py3
+from jedi._compatibility import is_py3, is_py35
 from token import *
 
 
@@ -24,6 +24,11 @@ else:
     tok_name[ELLIPSIS] = 'ELLIPSIS'
     N_TOKENS += 1
 
+if not is_py35:
+    ATEQUAL = N_TOKENS
+    tok_name[ATEQUAL] = 'ATEQUAL'
+    N_TOKENS += 1
+
 
 
 # Map from operator to number (since tokenize doesn't do this)
@@ -68,6 +73,7 @@ opmap_raw = """\
 %= PERCENTEQUAL
 &= AMPEREQUAL
 |= VBAREQUAL
+@= ATEQUAL
 ^= CIRCUMFLEXEQUAL
 <<= LEFTSHIFTEQUAL
 >>= RIGHTSHIFTEQUAL

jedi/parser/tokenize.py

@@ -76,7 +76,7 @@ triple = group("[uUbB]?[rR]?'''", '[uUbB]?[rR]?"""')
 # recognized as two instances of =).
 operator = group(r"\*\*=?", r">>=?", r"<<=?", r"!=",
                  r"//=?", r"->",
-                 r"[+\-*/%&|^=<>]=?",
+                 r"[+\-*@/%&|^=<>]=?",
                  r"~")
 
 bracket = '[][(){}]'

jedi/parser/tree.py

@@ -589,10 +589,44 @@ class BaseNode(Base):
         c = self.parent.children
         index = c.index(self)
         if index == len(c) - 1:
+            # TODO WTF? recursion?
             return self.get_next_leaf()
         else:
             return c[index + 1]
 
+    def last_leaf(self):
+        try:
+            return self.children[-1].last_leaf()
+        except AttributeError:
+            return self.children[-1]
+
+    def get_following_comment_same_line(self):
+        """
+        returns (as string) any comment that appears on the same line,
+        after the node, including the #
+        """
+        try:
+            if self.isinstance(ForStmt):
+                whitespace = self.children[5].first_leaf().prefix
+            elif self.isinstance(WithStmt):
+                whitespace = self.children[3].first_leaf().prefix
+            else:
+                whitespace = self.last_leaf().get_next_leaf().prefix
+        except AttributeError:
+            return None
+        except ValueError:
+            # TODO in some particular cases, the tree doesn't seem to be linked
+            # correctly
+            return None
+        if "#" not in whitespace:
+            return None
+        comment = whitespace[whitespace.index("#"):]
+        if "\r" in comment:
+            comment = comment[:comment.index("\r")]
+        if "\n" in comment:
+            comment = comment[:comment.index("\n")]
+        return comment
+
     @utf8_repr
     def __repr__(self):
         code = self.get_code().replace('\n', ' ').strip()
@@ -1471,7 +1505,7 @@ def _defined_names(current):
             names += _defined_names(child)
     elif is_node(current, 'atom'):
         names += _defined_names(current.children[1])
-    elif is_node(current, 'power'):
+    elif is_node(current, 'power', 'atom_expr'):
         if current.children[-2] != '**':  # Just if there's no operation
             trailer = current.children[-1]
             if trailer.children[0] == '.':