1  import difflib
   2  from pathlib import Path
   3  from typing import Dict, Iterable, Tuple
   4  
   5  from parso import split_lines
   6  
   7  from jedi.api.exceptions import RefactoringError
   8  from jedi.inference.value.namespace import ImplicitNSName
   9  
  10  EXPRESSION_PARTS = (
  11      'or_test and_test not_test comparison '
  12      'expr xor_expr and_expr shift_expr arith_expr term factor power atom_expr'
  13  ).split()
  14  
  15  
  16  class ChangedFile:
  17      def __init__(self, inference_state, from_path, to_path,
  18                   module_node, node_to_str_map):
  19          self._inference_state = inference_state
  20          self._from_path = from_path
  21          self._to_path = to_path
  22          self._module_node = module_node
  23          self._node_to_str_map = node_to_str_map
  24  
  25      def get_diff(self):
  26          old_lines = split_lines(self._module_node.get_code(), keepends=True)
  27          new_lines = split_lines(self.get_new_code(), keepends=True)
  28  
  29          # Add a newline at the end if it's missing. Otherwise the diff will be
  30          # very weird. A `diff -u file1 file2` would show the string:
  31          #
  32          #     \ No newline at end of file
  33          #
  34          # This is not necessary IMO, because Jedi does not really play with
  35          # newlines and the ending newline does not really matter in Python
  36          # files. ~dave
  37          if old_lines[-1] != '':
  38              old_lines[-1] += '\n'
  39          if new_lines[-1] != '':
  40              new_lines[-1] += '\n'
  41  
  42          project_path = self._inference_state.project.path
  43          if self._from_path is None:
  44              from_p = ''
  45          else:
  46              try:
  47                  from_p = self._from_path.relative_to(project_path)
  48              except ValueError:  # Happens it the path is not on th project_path
  49                  from_p = self._from_path
  50          if self._to_path is None:
  51              to_p = ''
  52          else:
  53              try:
  54                  to_p = self._to_path.relative_to(project_path)
  55              except ValueError:
  56                  to_p = self._to_path
  57          diff = difflib.unified_diff(
  58              old_lines, new_lines,
  59              fromfile=str(from_p),
  60              tofile=str(to_p),
  61          )
  62          # Apparently there's a space at the end of the diff - for whatever
  63          # reason.
  64          return ''.join(diff).rstrip(' ')
  65  
  66      def get_new_code(self):
  67          return self._inference_state.grammar.refactor(self._module_node, self._node_to_str_map)
  68  
  69      def apply(self):
  70          if self._from_path is None:
  71              raise RefactoringError(
  72                  'Cannot apply a refactoring on a Script with path=None'
  73              )
  74  
  75          with open(self._from_path, 'w', newline='') as f:
  76              f.write(self.get_new_code())
  77  
  78      def __repr__(self):
  79          return '<%s: %s>' % (self.__class__.__name__, self._from_path)
  80  
  81  
  82  class Refactoring:
  83      def __init__(self, inference_state, file_to_node_changes, renames=()):
  84          self._inference_state = inference_state
  85          self._renames = renames
  86          self._file_to_node_changes = file_to_node_changes
  87  
  88      def get_changed_files(self) -> Dict[Path, ChangedFile]:
  89          def calculate_to_path(p):
  90              if p is None:
  91                  return p
  92              p = str(p)
  93              for from_, to in renames:
  94                  if p.startswith(str(from_)):
  95                      p = str(to) + p[len(str(from_)):]
  96              return Path(p)
  97  
  98          renames = self.get_renames()
  99          return {
 100              path: ChangedFile(
 101                  self._inference_state,
 102                  from_path=path,
 103                  to_path=calculate_to_path(path),
 104                  module_node=next(iter(map_)).get_root_node(),
 105                  node_to_str_map=map_
 106              )
 107              # We need to use `or`, because the path can be None
 108              for path, map_ in sorted(
 109                  self._file_to_node_changes.items(),
 110                  key=lambda x: x[0] or Path("")
 111              )
 112          }
 113  
 114      def get_renames(self) -> Iterable[Tuple[Path, Path]]:
 115          """
 116          Files can be renamed in a refactoring.
 117          """
 118          return sorted(self._renames)
 119  
 120      def get_diff(self):
 121          text = ''
 122          project_path = self._inference_state.project.path
 123          for from_, to in self.get_renames():
 124              text += 'rename from %s\nrename to %s\n' \
 125                  % (_try_relative_to(from_, project_path), _try_relative_to(to, project_path))
 126  
 127          return text + ''.join(f.get_diff() for f in self.get_changed_files().values())
 128  
 129      def apply(self):
 130          """
 131          Applies the whole refactoring to the files, which includes renames.
 132          """
 133          for f in self.get_changed_files().values():
 134              f.apply()
 135  
 136          for old, new in self.get_renames():
 137              old.rename(new)
 138  
 139  
 140  def _calculate_rename(path, new_name):
 141      dir_ = path.parent
 142      if path.name in ('__init__.py', '__init__.pyi'):
 143          return dir_, dir_.parent.joinpath(new_name)
 144      return path, dir_.joinpath(new_name + path.suffix)
 145  
 146  
 147  def rename(inference_state, definitions, new_name):
 148      file_renames = set()
 149      file_tree_name_map = {}
 150  
 151      if not definitions:
 152          raise RefactoringError("There is no name under the cursor")
 153  
 154      for d in definitions:
 155          # This private access is ok in a way. It's not public to
 156          # protect Jedi users from seeing it.
 157          tree_name = d._name.tree_name
 158          if d.type == 'module' and tree_name is None and d.module_path is not None:
 159              p = Path(d.module_path)
 160              file_renames.add(_calculate_rename(p, new_name))
 161          elif isinstance(d._name, ImplicitNSName):
 162              for p in d._name._value.py__path__():
 163                  file_renames.add(_calculate_rename(Path(p), new_name))
 164          else:
 165              if tree_name is not None:
 166                  fmap = file_tree_name_map.setdefault(d.module_path, {})
 167                  fmap[tree_name] = tree_name.prefix + new_name
 168      return Refactoring(inference_state, file_tree_name_map, file_renames)
 169  
 170  
 171  def inline(inference_state, names):
 172      if not names:
 173          raise RefactoringError("There is no name under the cursor")
 174      if any(n.api_type in ('module', 'namespace') for n in names):
 175          raise RefactoringError("Cannot inline imports, modules or namespaces")
 176      if any(n.tree_name is None for n in names):
 177          raise RefactoringError("Cannot inline builtins/extensions")
 178  
 179      definitions = [n for n in names if n.tree_name.is_definition()]
 180      if len(definitions) == 0:
 181          raise RefactoringError("No definition found to inline")
 182      if len(definitions) > 1:
 183          raise RefactoringError("Cannot inline a name with multiple definitions")
 184      if len(names) == 1:
 185          raise RefactoringError("There are no references to this name")
 186  
 187      tree_name = definitions[0].tree_name
 188  
 189      expr_stmt = tree_name.get_definition()
 190      if expr_stmt.type != 'expr_stmt':
 191          type_ = dict(
 192              funcdef='function',
 193              classdef='class',
 194          ).get(expr_stmt.type, expr_stmt.type)
 195          raise RefactoringError("Cannot inline a %s" % type_)
 196  
 197      if len(expr_stmt.get_defined_names(include_setitem=True)) > 1:
 198          raise RefactoringError("Cannot inline a statement with multiple definitions")
 199      first_child = expr_stmt.children[1]
 200      if first_child.type == 'annassign' and len(first_child.children) == 4:
 201          first_child = first_child.children[2]
 202      if first_child != '=':
 203          if first_child.type == 'annassign':
 204              raise RefactoringError(
 205                  'Cannot inline a statement that is defined by an annotation'
 206              )
 207          else:
 208              raise RefactoringError(
 209                  'Cannot inline a statement with "%s"'
 210                  % first_child.get_code(include_prefix=False)
 211              )
 212  
 213      rhs = expr_stmt.get_rhs()
 214      replace_code = rhs.get_code(include_prefix=False)
 215  
 216      references = [n for n in names if not n.tree_name.is_definition()]
 217      file_to_node_changes = {}
 218      for name in references:
 219          tree_name = name.tree_name
 220          path = name.get_root_context().py__file__()
 221          s = replace_code
 222          if rhs.type == 'testlist_star_expr' \
 223                  or tree_name.parent.type in EXPRESSION_PARTS \
 224                  or tree_name.parent.type == 'trailer' \
 225                  and tree_name.parent.get_next_sibling() is not None:
 226              s = '(' + replace_code + ')'
 227  
 228          of_path = file_to_node_changes.setdefault(path, {})
 229  
 230          n = tree_name
 231          prefix = n.prefix
 232          par = n.parent
 233          if par.type == 'trailer' and par.children[0] == '.':
 234              prefix = par.parent.children[0].prefix
 235              n = par
 236              for some_node in par.parent.children[:par.parent.children.index(par)]:
 237                  of_path[some_node] = ''
 238          of_path[n] = prefix + s
 239  
 240      path = definitions[0].get_root_context().py__file__()
 241      changes = file_to_node_changes.setdefault(path, {})
 242      changes[expr_stmt] = _remove_indent_of_prefix(expr_stmt.get_first_leaf().prefix)
 243      next_leaf = expr_stmt.get_next_leaf()
 244  
 245      # Most of the time we have to remove the newline at the end of the
 246      # statement, but if there's a comment we might not need to.
 247      if next_leaf.prefix.strip(' \t') == '' \
 248              and (next_leaf.type == 'newline' or next_leaf == ';'):
 249          changes[next_leaf] = ''
 250      return Refactoring(inference_state, file_to_node_changes)
 251  
 252  
 253  def _remove_indent_of_prefix(prefix):
 254      r"""
 255      Removes the last indentation of a prefix, e.g. " \n \n " becomes " \n \n".
 256      """
 257      return ''.join(split_lines(prefix, keepends=True)[:-1])
 258  
 259  
 260  def _try_relative_to(path: Path, base: Path) -> Path:
 261      try:
 262          return path.relative_to(base)
 263      except ValueError:
 264          return path
 265