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