1# This file is part of the lib3to6 project 

2# https://github.com/mbarkhau/lib3to6 

3# 

4# Copyright (c) 2019-2021 Manuel Barkhau (mbarkhau@gmail.com) - MIT License 

5# SPDX-License-Identifier: MIT 

6 

7import ast 

8import typing as typ 

9 

10from . import common 

11from . import fixer_base as fb 

12 

13 

14class NamedExprFixer(fb.TransformerFixerBase): 

15 

16 version_info = common.VersionInfo(apply_since="2.7", apply_until="3.7") 

17 

18 def _extract_and_replace_named_exprs( 

19 self, expr: ast.expr 

20 ) -> typ.Tuple[typ.List[ast.stmt], ast.expr]: 

21 new_assigns: typ.List[ast.stmt] = [] 

22 if isinstance(expr, ast.NamedExpr): 

23 new_assigns.append(ast.Assign(targets=[expr.target], value=expr.value)) 

24 new_expr = ast.Name(id=expr.target.id, ctx=ast.Load()) 

25 return new_assigns, new_expr 

26 else: 

27 if isinstance(expr, ast.UnaryOp): 

28 new_sub_assigns, new_operand = self._extract_and_replace_named_exprs(expr.operand) 

29 new_assigns.extend(new_sub_assigns) 

30 expr.operand = new_operand 

31 if isinstance(expr, (ast.BinOp, ast.Compare)): 

32 new_sub_assigns, new_left = self._extract_and_replace_named_exprs(expr.left) 

33 new_assigns.extend(new_sub_assigns) 

34 expr.left = new_left 

35 if isinstance(expr, ast.BinOp): 

36 new_sub_assigns, new_right = self._extract_and_replace_named_exprs(expr.right) 

37 new_assigns.extend(new_sub_assigns) 

38 expr.right = new_right 

39 if isinstance(expr, ast.BoolOp): 

40 new_values = [] 

41 for comp in expr.values: 

42 new_sub_assigns, new_comp = self._extract_and_replace_named_exprs(comp) 

43 new_values.append(new_comp) 

44 new_assigns.extend(new_sub_assigns) 

45 expr.values = new_values 

46 if isinstance(expr, ast.Compare): 

47 new_comparators = [] 

48 for comp in expr.comparators: 

49 new_sub_assigns, new_comp = self._extract_and_replace_named_exprs(comp) 

50 new_comparators.append(new_comp) 

51 new_assigns.extend(new_sub_assigns) 

52 expr.comparators = new_comparators 

53 

54 return new_assigns, expr 

55 

56 def _update(self, nodelist: typ.List[ast.stmt], indent: int = 0) -> None: 

57 i = 0 

58 while i < len(nodelist): 

59 node = nodelist[i] 

60 if isinstance(node, (ast.If, ast.While)): 

61 new_assigns, new_test = self._extract_and_replace_named_exprs(node.test) 

62 if new_assigns and isinstance(node, ast.While): 

63 loopcond_name = '__loop_condition' 

64 # __loop_condition = True 

65 loopcond_init_node = ast.Assign( 

66 targets=[ast.Name(id=loopcond_name, ctx=ast.Store())], 

67 value=ast.NameConstant(value=True, kind=None), 

68 ) 

69 nodelist.insert(i, loopcond_init_node) 

70 

71 # while __loop_condition: 

72 node.test = ast.Name(id=loopcond_name, ctx=ast.Load()) 

73 

74 # __loop_condition = test 

75 loopcond_assign_node = ast.Assign( 

76 targets=[ast.Name(id=loopcond_name, ctx=ast.Store())], 

77 value=new_test, 

78 ) 

79 # if __loop_condition: 

80 new_ifnode = ast.If( 

81 test=ast.Name(id=loopcond_name, ctx=ast.Load()), 

82 body=node.body, 

83 orelse=[], 

84 ) 

85 node.body = new_assigns + [loopcond_assign_node, new_ifnode] 

86 i += 1 

87 else: 

88 node.test = new_test 

89 if new_assigns: 

90 nodelist[i:i] = new_assigns 

91 i += 1 + len(new_assigns) 

92 else: 

93 i += 1 

94 

95 for nodelist_name in ('body', 'orelse', 'handlers', 'finalbody'): 

96 sub_nodelist = getattr(node, nodelist_name, None) 

97 if sub_nodelist: 

98 self._update(sub_nodelist, indent + 1) 

99 

100 def apply_fix(self, ctx: common.BuildContext, tree: ast.Module) -> ast.Module: 

101 self._update(tree.body) 

102 return tree