Skip to content

Commit

Permalink
fix: some fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
15r10nk committed Aug 6, 2023
1 parent 67febea commit 454b71b
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 14 deletions.
4 changes: 2 additions & 2 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,11 @@ repos:
rev: v3.9.0
hooks:
- args:
- --py38-plus
- --py37-plus
id: reorder-python-imports
- hooks:
- args:
- --py38-plus
- --py37-plus
id: pyupgrade
repo: https://github.com/asottile/pyupgrade
rev: v3.3.1
Expand Down
36 changes: 24 additions & 12 deletions pysource_minimize/_minimize.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
try:
from ast import unparse
except ImportError:

from astunparse import unparse # type: ignore

TESTING = "PYTEST_CURRENT_TEST" in os.environ
Expand Down Expand Up @@ -36,7 +35,10 @@ def __init__(self, source, checker, progress_callback):
self.progress_callback = progress_callback

self.original_source = source
self.original_ast = ast.parse(source)
if sys.version_info >= (3, 8):
self.original_ast = ast.parse(source, type_comments=True)
else:
self.original_ast = ast.parse(source)

class UniqueObj(ast.NodeTransformer):
def visit(self, node):
Expand Down Expand Up @@ -182,8 +184,7 @@ def try_only(self, node, child) -> bool:

def minimize_comprehension(self, comp):
self.minimize_expr(comp.iter)
for if_ in comp.ifs:
self.minimize_expr(if_)
self.minimize_list(comp.ifs, terminal=self.minimize_expr)

def try_only_minimize(self, node, *childs):
childs = [child for child in childs if child is not None]
Expand Down Expand Up @@ -277,14 +278,6 @@ def minimize_expr(self, node):
)
return

not_stared = [arg for arg in node.args if not isinstance(arg, ast.Starred)]
for arg in not_stared:
if self.try_only(node, arg):
self.minimize(arg)
break
else:
self.minimize(node.args)

elif isinstance(
node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)
):
Expand Down Expand Up @@ -312,6 +305,8 @@ def minimize_expr(self, node):
self.minimize_expr(node.elt)
self.minimize_list(node.generators, self.minimize_comprehension, 1)

elif isinstance(node, ast.NamedExpr):
self.try_only_minimize(node, node.value)
else:
raise TypeError(node) # Expr

Expand Down Expand Up @@ -353,6 +348,8 @@ def minimize_stmt(self, node):
self.minimize_list(nargs.kwonlyargs, lambda e: None)
self.try_none(nargs.vararg)
self.try_none(nargs.kwarg)
if node.returns:
self.minimize_expr(node.returns)

elif isinstance(node, ast.ClassDef):
if self.try_only_minimize(node, node.decorator_list):
Expand All @@ -361,7 +358,22 @@ def minimize_stmt(self, node):
if self.try_only_minimize(node, node.body):
return

for e in [
*[kw.value for kw in node.keywords],
*[
arg.value if isinstance(arg, ast.Starred) else arg
for arg in node.bases
],
]:
if self.try_only(node, e):
self.minimize(e)
return

self.minimize(node.bases)
self.minimize_list(
node.keywords, terminal=lambda kw: self.minimize(kw.value)
)
return

elif isinstance(node, ast.Return):
self.try_only_minimize(node, node.value)
Expand Down

0 comments on commit 454b71b

Please sign in to comment.