1 | #pragma once |
---|---|
2 | |
3 | #include <ATen/core/dynamic_type.h> |
4 | #include <ATen/core/jit_type.h> |
5 | |
6 | namespace c10 { |
7 | |
8 | class TORCH_API TypeParser { |
9 | public: |
10 | explicit TypeParser(std::string pythonStr); |
11 | explicit TypeParser(std::vector<std::string>& pythonStrs); |
12 | |
13 | TypePtr parse(); |
14 | std::vector<TypePtr> parseList(); |
15 | static const std::unordered_set<std::string>& getNonSimpleType(); |
16 | static const std::unordered_set<std::string>& getCustomType(); |
17 | std::unordered_set<std::string> getContainedTypes(); |
18 | |
19 | private: |
20 | TypePtr parseNamedTuple(const std::string& qualified_name); |
21 | TypePtr parseCustomType(); |
22 | TypePtr parseTorchbindClassType(); |
23 | TypePtr parseNonSimple(const std::string& token); |
24 | |
25 | void expect(const char* s); |
26 | void expectChar(char c); |
27 | template <typename T> |
28 | TypePtr parseSingleElementType(); |
29 | |
30 | void lex(); |
31 | |
32 | std::string next(); |
33 | c10::string_view nextView(); |
34 | void advance(); |
35 | C10_NODISCARD c10::string_view cur() const; |
36 | |
37 | std::string pythonStr_; |
38 | size_t start_; |
39 | c10::string_view next_token_; |
40 | |
41 | // Used for parsing string list |
42 | std::vector<std::string> pythonStrs_; |
43 | std::unordered_map<std::string, c10::TypePtr> str_type_ptr_map_; |
44 | |
45 | // Store all contained types when parsing a string |
46 | std::unordered_set<std::string> contained_types_; |
47 | }; |
48 | |
49 | TORCH_API TypePtr parseType(const std::string& pythonStr); |
50 | |
51 | TORCH_API std::vector<TypePtr> parseType(std::vector<std::string>& pythonStr); |
52 | |
53 | } // namespace c10 |
54 |