Skip to content

Commit

Permalink
[JIT] Resolve string literal type annotations using `Resolver::resolv…
Browse files Browse the repository at this point in the history
…eType`

**Summary**
This commit modifies `ScriptTypeParser::parseTypeFromExpr` so that
string literal type annotations are resolved using
`Resolver::resolveType`. At present, they are parsed in
`parseBaseTypeName`, which inadvertently allows any key from
`string_to_type_lut` to be used as a string literal type annotation.

**Test Plan**
Existing unit tests (most notably
`TestClassType.test_self_referential_method` which tests the main
feature, self-referential class type annotations, that make use of
string literal type annotations).

**Fixes**
This commit fixes #47570.

ghstack-source-id: 9e5a2c83f6211084a378df1863e662e2ced936ce
Pull Request resolved: #47731
  • Loading branch information
Meghan Lele committed Nov 11, 2020
1 parent 0a7ebf0 commit 8f51c5c
Showing 1 changed file with 9 additions and 3 deletions.
12 changes: 9 additions & 3 deletions torch/csrc/jit/frontend/script_type_parser.cpp
Expand Up @@ -149,9 +149,6 @@ c10::optional<std::string> ScriptTypeParser::parseBaseTypeName(
case TK_NONE: {
return "None";
}
case TK_STRINGLITERAL: {
return StringLiteral(expr).text();
}
case '.': {
auto select = Select(expr);
const std::string& name = select.selector().name();
Expand Down Expand Up @@ -190,6 +187,15 @@ TypePtr ScriptTypeParser::parseTypeFromExprImpl(const Expr& expr) const {
}
return subscriptToType(*value_name, subscript);

} else if (expr.kind() == TK_STRINGLITERAL) {
auto type_name = StringLiteral(expr).text();
if (resolver_) {
if (auto typePtr = resolver_->resolveType(type_name, expr.range())) {
return typePtr;
}
}

throw ErrorReport(expr) << "Unknown type name '" << type_name << "'";
} else if (auto name = parseBaseTypeName(expr)) {
auto itr = string_to_type_lut().find(*name);
if (itr != string_to_type_lut().end()) {
Expand Down

0 comments on commit 8f51c5c

Please sign in to comment.