Commit d619a731 authored by Francis Herne's avatar Francis Herne

Improve support for 'with' statements.

The previous code didn't look at `__enter__()`, and assumed that
 all context-managers returned their own type.

We also didn't account for targets other than simple names,
 e.g. `with Mgr() as (foo, bar):`

Thanks to Nicolás Alvarez.

BUG: 399533
BUG: 399534

Differential Revision: https://phabricator.kde.org/D16085
parent d169b27e
......@@ -152,6 +152,7 @@ class _io_TextIOWrapper():
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 = ""
......
......@@ -364,7 +364,21 @@ void DeclarationBuilder::visitWithItem(WithItemAst* node)
// 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<Declaration>(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<FunctionDeclaration*>(
Helper::accessAttribute(mgrType, enterId, topContext()))) {
if ( auto enterFuncType = enterFunc->type<FunctionType>() ) {
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);
}
......
......@@ -871,7 +871,25 @@ void PyDUChainTest::testTypes_data()
" 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";
......
......@@ -368,7 +368,7 @@ class KDEVPYTHONPARSER_EXPORT WithItemAst : public Ast {
public:
WithItemAst(Ast* parent);
ExpressionAst* contextExpression;
NameAst* optionalVars;
ExpressionAst* optionalVars;
};
class KDEVPYTHONPARSER_EXPORT WithAst : public StatementAst {
......
......@@ -780,7 +780,7 @@ private:
if ( ! node ) return nullptr;
WithItemAst* v = new WithItemAst(parent());
nodeStack.push(v); v->contextExpression = static_cast<ExpressionAst*>(visitNode(node->context_expr)); nodeStack.pop();
nodeStack.push(v); v->optionalVars = static_cast<NameAst*>(visitNode(node->optional_vars)); nodeStack.pop();
nodeStack.push(v); v->optionalVars = static_cast<ExpressionAst*>(visitNode(node->optional_vars)); nodeStack.pop();
return v;
}
......
......@@ -99,4 +99,4 @@ RULE_FOR _arguments;KIND any;ACTIONS create|ArgumentsAst set|vararg->ArgAst,vara
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;;
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment