frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-11 22:34:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

[JIT] Make NoneType annotation_str emit NoneType instead of None (#54642)

Browse Source 
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/54642

Test Plan: Imported from OSS

Reviewed By: SplitInfinity

Differential Revision: D27314174

Pulled By: jamesr66a

fbshipit-source-id: 153e9aa4ab781fa1d49d9d55a2e487bf7b04f0d7
This commit is contained in:
James Reed
2021-03-26 11:28:42 -07:00
committed by Facebook GitHub Bot
parent 1e9ad6e5cd
commit 3db2333d09
7 changed files with 20 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
test/test_jit.py
View File

@ -11562,6 +11562,15 @@ dedent """
return x
''')
def test_parse_none_type_annotation(self):
cu = torch.jit.CompilationUnit('''
def foo(x : NoneType) -> NoneType:
return x
''')
foo_code = cu.find_function('foo').code
FileCheck().check(": None").check("-> None").run(foo_code)
def test_zip_enumerate_modulelist(self):
class Sub(torch.nn.Module):
def __init__(self):

3
torch/csrc/jit/frontend/lexer.h
View File

@ -119,7 +119,8 @@ namespace jit {
_(TK_WITH_ITEM, "withitem", "") \
_(TK_AS, "as", "as") \
_(TK_PROP, "property", "") \
_(TK_ELLIPSIS, "Ellipsis", "Ellipsis")
_(TK_ELLIPSIS, "Ellipsis", "Ellipsis") \
_(TK_NONE_TYPE, "NoneType", "NoneType")
enum TokenKind {
// we use characters to represent themselves so skip all valid characters

3
torch/csrc/jit/frontend/parser.cpp
View File

@ -115,7 +115,8 @@ struct ParserImpl {
} break;
case TK_TRUE:
case TK_FALSE:
case TK_NONE: {
case TK_NONE:
case TK_NONE_TYPE: {
auto k = L.cur().kind;
auto r = L.cur().range;
prefix = create_compound(k, r, {});

3
torch/csrc/jit/frontend/schema_type_parser.cpp
View File

@ -60,13 +60,14 @@ TypePtr SchemaTypeParser::parseBaseType() {
{"int", IntType::get()},
{"bool", BoolType::get()},
{"None", NoneType::get()},
{"NoneType", NoneType::get()},
{"Capsule", CapsuleType::get()},
{"Any", at::AnyType::get()},
{"AnyClassType", at::AnyClassType::get()},
{"AnyEnumType", at::AnyEnumType::get()},
};
auto tok = L.cur();
if (!L.nextIf(TK_NONE)) {
if (!L.nextIf(TK_NONE) && !L.nextIf(TK_NONE_TYPE)) {
L.expect(TK_IDENT);
}
std::string text = tok.text();

3
torch/csrc/jit/frontend/script_type_parser.cpp
View File

@ -203,6 +203,9 @@ c10::optional<std::string> ScriptTypeParser::parseBaseTypeName(
case TK_NONE: {
return "None";
}
case TK_NONE_TYPE: {
return "NoneType";
}
case '.': {
auto select = Select(expr);
const std::string& name = select.selector().name();

1
torch/csrc/jit/frontend/string_to_type.cpp
View File

@ -17,6 +17,7 @@ const std::unordered_map<std::string, TypePtr>& string_to_type_lut() {
// parsing serialized methods that use implicit conversions to Scalar
{"number", NumberType::get()},
{"None", NoneType::get()},
{"NoneType", NoneType::get()},
{"Any", AnyType::get()},
{"Capsule", CapsuleType::get()},
{"list", AnyListType::get()},

1
torch/csrc/jit/frontend/tree_views.h
View File

@ -292,6 +292,7 @@ struct Expr : public TreeView {
case TK_TRUE:
case TK_FALSE:
case TK_NONE:
case TK_NONE_TYPE:
case TK_CAST:
case TK_APPLY:
case '.':