diff --git a/documentation_files/builtindocumentation.py b/documentation_files/builtindocumentation.py --- a/documentation_files/builtindocumentation.py +++ b/documentation_files/builtindocumentation.py @@ -152,6 +152,7 @@ def writelines(self,sequence): return None def __iter__(self): return self def __next__(self): return "" + def __enter__(self): return self buffer = _io_TextIOWrapper() # Not quite closed = True encoding = "" diff --git a/duchain/declarationbuilder.cpp b/duchain/declarationbuilder.cpp --- a/duchain/declarationbuilder.cpp +++ b/duchain/declarationbuilder.cpp @@ -364,7 +364,21 @@ // For statements like "with open(f) as x", a new variable must be created; do this here. ExpressionVisitor v(currentContext()); v.visitNode(node->contextExpression); - visitVariableDeclaration(node->optionalVars, nullptr, v.lastType()); + auto mgrType = v.lastType(); + auto enterType = mgrType; // If we can't find __enter__(), assume it returns `self` like file objects. + + static const IndexedIdentifier enterId(KDevelop::Identifier("__enter__")); + + DUChainReadLocker lock; + if ( auto enterFunc = dynamic_cast( + Helper::accessAttribute(mgrType, enterId, topContext()))) { + if ( auto enterFuncType = enterFunc->type() ) { + enterType = enterFuncType->returnType(); + } + } + lock.unlock(); + // This may be any assignable expression, e.g. `with foo() as bar[3]: ...` + assignToUnknown(node->optionalVars, enterType); } Python::AstDefaultVisitor::visitWithItem(node); } diff --git a/duchain/tests/pyduchaintest.cpp b/duchain/tests/pyduchaintest.cpp --- a/duchain/tests/pyduchaintest.cpp +++ b/duchain/tests/pyduchaintest.cpp @@ -871,7 +871,25 @@ " return x\n" "f = Foo()\n" "checkme = f.func()" << "int"; + QTest::newRow("with") << "with open('foo') as f: checkme = f.read()" << "str"; + QTest::newRow("with_list_target") << "bar = [1, 2, 3]\n" + "with open('foo') as bar[1]: checkme = bar[1].read()" << "str"; + QTest::newRow("with_attr_target") << "bar = object()\n" + "with open('foo') as bar.zep: checkme = bar.zep.read()" << "str"; + QTest::newRow("with_nonself_enter") << // From https://bugs.kde.org/show_bug.cgi?id=399534 + "class Mgr:\n" + " def __enter__(self): return 42\n" + " def __exit__(self, *args): pass\n" + "with Mgr() as asd:\n" + " checkme = asd" << "int"; + QTest::newRow("with_tuple_target") << + "class Mgr:\n" + " def __enter__(self): return (42, 3.4)\n" + " def __exit__(self, *args): pass\n" + "with Mgr() as (aa, bb):\n" + " checkme = bb" << "float"; + QTest::newRow("arg_after_vararg") << "def func(x, y, *, z:int): return z\ncheckme = func()" << "int"; QTest::newRow("arg_after_vararg_with_default") << "def func(x=5, y=3, *, z:int): return z\ncheckme = func()" << "int"; diff --git a/parser/ast.h b/parser/ast.h --- a/parser/ast.h +++ b/parser/ast.h @@ -368,7 +368,7 @@ public: WithItemAst(Ast* parent); ExpressionAst* contextExpression; - NameAst* optionalVars; + ExpressionAst* optionalVars; }; class KDEVPYTHONPARSER_EXPORT WithAst : public StatementAst { diff --git a/parser/generated.h b/parser/generated.h --- a/parser/generated.h +++ b/parser/generated.h @@ -780,7 +780,7 @@ if ( ! node ) return nullptr; WithItemAst* v = new WithItemAst(parent()); nodeStack.push(v); v->contextExpression = static_cast(visitNode(node->context_expr)); nodeStack.pop(); - nodeStack.push(v); v->optionalVars = static_cast(visitNode(node->optional_vars)); nodeStack.pop(); + nodeStack.push(v); v->optionalVars = static_cast(visitNode(node->optional_vars)); nodeStack.pop(); return v; } diff --git a/parser/python36.sdef b/parser/python36.sdef --- a/parser/python36.sdef +++ b/parser/python36.sdef @@ -99,4 +99,4 @@ RULE_FOR _arg;KIND any;ACTIONS create|ArgAst set|argumentName~>arg set|annotation->ExpressionAst,annotation;; RULE_FOR _keyword;KIND any;ACTIONS create|KeywordAst set|argumentName~>arg set|value->ExpressionAst,value;; RULE_FOR _alias;KIND any;ACTIONS create|AliasAst set|name~>name set|asName~>asname;; -RULE_FOR _withitem;KIND any; ACTIONS create|WithItemAst set|contextExpression->ExpressionAst,context_expr set|optionalVars->NameAst,optional_vars;; +RULE_FOR _withitem;KIND any; ACTIONS create|WithItemAst set|contextExpression->ExpressionAst,context_expr set|optionalVars->ExpressionAst,optional_vars;;