1
- from typing import List , Tuple
1
+ from abc import abstractmethod , ABC
2
+ from typing import List , Tuple , Optional , Protocol
2
3
3
4
4
5
class ImportStatement :
@@ -64,17 +65,22 @@ def accept(self, visitor):
64
65
return visitor .dispatch (self )
65
66
66
67
67
- class ImportInfo :
68
- def get_imported_primaries (self , context ):
69
- pass
68
+ class ImportInfo (ABC ):
69
+ names_and_aliases : List [Tuple [str , Optional [str ]]]
70
+
71
+ @abstractmethod
72
+ def get_imported_primaries (self , context ) -> List [str ]: ...
70
73
71
74
def get_imported_names (self , context ):
72
75
return [
73
76
primary .split ("." )[0 ] for primary in self .get_imported_primaries (context )
74
77
]
75
78
76
- def get_import_statement (self ):
77
- pass
79
+ def __repr__ (self ):
80
+ return f'<{ self .__class__ .__name__ } "{ self .get_import_statement ()} ">'
81
+
82
+ @abstractmethod
83
+ def get_import_statement (self ) -> str : ...
78
84
79
85
def is_empty (self ):
80
86
pass
@@ -105,10 +111,13 @@ def get_empty_import():
105
111
106
112
107
113
class NormalImport (ImportInfo ):
108
- def __init__ (self , names_and_aliases ):
114
+ def __init__ (
115
+ self ,
116
+ names_and_aliases : List [Tuple [str , Optional [str ]]],
117
+ ) -> None :
109
118
self .names_and_aliases = names_and_aliases
110
119
111
- def get_imported_primaries (self , context ):
120
+ def get_imported_primaries (self , context ) -> List [ str ] :
112
121
result = []
113
122
for name , alias in self .names_and_aliases :
114
123
if alias :
@@ -117,7 +126,7 @@ def get_imported_primaries(self, context):
117
126
result .append (name )
118
127
return result
119
128
120
- def get_import_statement (self ):
129
+ def get_import_statement (self ) -> str :
121
130
result = "import "
122
131
for name , alias in self .names_and_aliases :
123
132
result += name
@@ -131,12 +140,20 @@ def is_empty(self):
131
140
132
141
133
142
class FromImport (ImportInfo ):
134
- def __init__ (self , module_name , level , names_and_aliases ):
143
+ module_name : str
144
+ level : int
145
+
146
+ def __init__ (
147
+ self ,
148
+ module_name : str ,
149
+ level : int ,
150
+ names_and_aliases : List [Tuple [str , Optional [str ]]],
151
+ ):
135
152
self .module_name = module_name
136
153
self .level = level
137
154
self .names_and_aliases = names_and_aliases
138
155
139
- def get_imported_primaries (self , context ):
156
+ def get_imported_primaries (self , context ) -> List [ str ] :
140
157
if self .names_and_aliases [0 ][0 ] == "*" :
141
158
module = self .get_imported_module (context )
142
159
return [name for name in module if not name .startswith ("_" )]
@@ -173,7 +190,7 @@ def get_imported_module(self, context):
173
190
self .module_name , context .folder , self .level
174
191
)
175
192
176
- def get_import_statement (self ):
193
+ def get_import_statement (self ) -> str :
177
194
result = "from " + "." * self .level + self .module_name + " import "
178
195
for name , alias in self .names_and_aliases :
179
196
result += name
@@ -190,14 +207,17 @@ def is_star_import(self):
190
207
191
208
192
209
class EmptyImport (ImportInfo ):
193
- names_and_aliases : List [ Tuple [ str , str ]] = []
210
+ names_and_aliases = []
194
211
195
212
def is_empty (self ):
196
213
return True
197
214
198
- def get_imported_primaries (self , context ):
215
+ def get_imported_primaries (self , context ) -> List [ str ] :
199
216
return []
200
217
218
+ def get_import_statement (self ) -> str :
219
+ raise NotImplementedError ()
220
+
201
221
202
222
class ImportContext :
203
223
def __init__ (self , project , folder ):
0 commit comments