@@ -58,18 +58,18 @@ def eq_sub_func(match: re.Match) -> str:
5858 else :
5959 return "\n " .join (
6060 [
61- f"{ match .group ('indent' )} def __eq__(self, arg0 : { match .group ('other' )} ) -> { match .group ('return' )} :{ match .group ('ellipsis_docstring' )} " ,
61+ f"{ match .group ('indent' )} def __eq__(self, other : { match .group ('other' )} ) -> { match .group ('return' )} :{ match .group ('ellipsis_docstring' )} " ,
6262 f"{ match .group ('indent' )} @typing.overload" ,
63- f"{ match .group ('indent' )} def __eq__(self, arg0 : typing.Any) -> bool | types.NotImplementedType: ..." ,
63+ f"{ match .group ('indent' )} def __eq__(self, other : typing.Any) -> bool | types.NotImplementedType: ..." ,
6464 ]
6565 )
6666 else :
6767 return "\n " .join (
6868 [
6969 f"{ match .group ('indent' )} @typing.overload" ,
70- f"{ match .group ('indent' )} def __eq__(self, arg0 : { match .group ('other' )} ) -> { match .group ('return' )} :{ match .group ('ellipsis_docstring' )} " ,
70+ f"{ match .group ('indent' )} def __eq__(self, other : { match .group ('other' )} ) -> { match .group ('return' )} :{ match .group ('ellipsis_docstring' )} " ,
7171 f"{ match .group ('indent' )} @typing.overload" ,
72- f"{ match .group ('indent' )} def __eq__(self, arg0 : typing.Any) -> bool | types.NotImplementedType: ..." ,
72+ f"{ match .group ('indent' )} def __eq__(self, other : typing.Any) -> bool | types.NotImplementedType: ..." ,
7373 ]
7474 )
7575
@@ -246,13 +246,12 @@ def main() -> None:
246246 pyi = EqPattern .sub (eq_sub_func , pyi )
247247 pyi = pyi .replace ("**kwargs)" , "**kwargs: typing.Any)" )
248248 pyi_split = [l .rstrip ("\r " ) for l in pyi .split ("\n " )]
249- for hidden_import in ["amulet.nbt" ]:
249+ for hidden_import in ["amulet.nbt" , "typing" , "types" ]:
250250 if hidden_import in pyi and f"import { hidden_import } " not in pyi_split :
251- pyi_split .insert (2 , f"import { hidden_import } " )
252- if "import typing" not in pyi_split :
253- pyi_split .insert (2 , "import typing" )
254- if "import types" not in pyi_split :
255- pyi_split .insert (2 , "import types" )
251+ pyi_split .insert (
252+ pyi_split .index ("from __future__ import annotations" ) + 1 ,
253+ f"import { hidden_import } " ,
254+ )
256255 pyi = "\n " .join (pyi_split )
257256 with open (stub_path , "w" , encoding = "utf-8" ) as f :
258257 f .write (pyi )
0 commit comments