1- from typing import List , Tuple
1+ from abc import abstractmethod , ABC
2+ from typing import List , Tuple , Optional , Protocol
23
34
45class ImportStatement :
@@ -64,17 +65,22 @@ def accept(self, visitor):
6465 return visitor .dispatch (self )
6566
6667
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 ]: ...
7073
7174 def get_imported_names (self , context ):
7275 return [
7376 primary .split ("." )[0 ] for primary in self .get_imported_primaries (context )
7477 ]
7578
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 : ...
7884
7985 def is_empty (self ):
8086 pass
@@ -105,10 +111,13 @@ def get_empty_import():
105111
106112
107113class 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 :
109118 self .names_and_aliases = names_and_aliases
110119
111- def get_imported_primaries (self , context ):
120+ def get_imported_primaries (self , context ) -> List [ str ] :
112121 result = []
113122 for name , alias in self .names_and_aliases :
114123 if alias :
@@ -117,7 +126,7 @@ def get_imported_primaries(self, context):
117126 result .append (name )
118127 return result
119128
120- def get_import_statement (self ):
129+ def get_import_statement (self ) -> str :
121130 result = "import "
122131 for name , alias in self .names_and_aliases :
123132 result += name
@@ -131,12 +140,20 @@ def is_empty(self):
131140
132141
133142class 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+ ):
135152 self .module_name = module_name
136153 self .level = level
137154 self .names_and_aliases = names_and_aliases
138155
139- def get_imported_primaries (self , context ):
156+ def get_imported_primaries (self , context ) -> List [ str ] :
140157 if self .names_and_aliases [0 ][0 ] == "*" :
141158 module = self .get_imported_module (context )
142159 return [name for name in module if not name .startswith ("_" )]
@@ -173,7 +190,7 @@ def get_imported_module(self, context):
173190 self .module_name , context .folder , self .level
174191 )
175192
176- def get_import_statement (self ):
193+ def get_import_statement (self ) -> str :
177194 result = "from " + "." * self .level + self .module_name + " import "
178195 for name , alias in self .names_and_aliases :
179196 result += name
@@ -190,14 +207,17 @@ def is_star_import(self):
190207
191208
192209class EmptyImport (ImportInfo ):
193- names_and_aliases : List [ Tuple [ str , str ]] = []
210+ names_and_aliases = []
194211
195212 def is_empty (self ):
196213 return True
197214
198- def get_imported_primaries (self , context ):
215+ def get_imported_primaries (self , context ) -> List [ str ] :
199216 return []
200217
218+ def get_import_statement (self ) -> str :
219+ raise NotImplementedError ()
220+
201221
202222class ImportContext :
203223 def __init__ (self , project , folder ):
0 commit comments