1  import re
   2  import textwrap
   3  from ast import literal_eval
   4  from inspect import cleandoc
   5  from weakref import WeakKeyDictionary
   6  
   7  from parso.python import tree
   8  from parso.cache import parser_cache
   9  from parso import split_lines
  10  
  11  _EXECUTE_NODES = {'funcdef', 'classdef', 'import_from', 'import_name', 'test',
  12                    'or_test', 'and_test', 'not_test', 'comparison', 'expr',
  13                    'xor_expr', 'and_expr', 'shift_expr', 'arith_expr',
  14                    'atom_expr', 'term', 'factor', 'power', 'atom'}
  15  
  16  _FLOW_KEYWORDS = (
  17      'try', 'except', 'finally', 'else', 'if', 'elif', 'with', 'for', 'while'
  18  )
  19  
  20  
  21  def get_executable_nodes(node, last_added=False):
  22      """
  23      For static analysis.
  24      """
  25      result = []
  26      typ = node.type
  27      if typ == 'name':
  28          next_leaf = node.get_next_leaf()
  29          if last_added is False and node.parent.type != 'param' and next_leaf != '=':
  30              result.append(node)
  31      elif typ == 'expr_stmt':
  32          # I think inferring the statement (and possibly returned arrays),
  33          # should be enough for static analysis.
  34          result.append(node)
  35          for child in node.children:
  36              result += get_executable_nodes(child, last_added=True)
  37      elif typ == 'decorator':
  38          # decorator
  39          if node.children[-2] == ')':
  40              node = node.children[-3]
  41              if node != '(':
  42                  result += get_executable_nodes(node)
  43      else:
  44          try:
  45              children = node.children
  46          except AttributeError:
  47              pass
  48          else:
  49              if node.type in _EXECUTE_NODES and not last_added:
  50                  result.append(node)
  51  
  52              for child in children:
  53                  result += get_executable_nodes(child, last_added)
  54  
  55      return result
  56  
  57  
  58  def get_sync_comp_fors(comp_for):
  59      yield comp_for
  60      last = comp_for.children[-1]
  61      while True:
  62          if last.type == 'comp_for':
  63              yield last.children[1]  # Ignore the async.
  64          elif last.type == 'sync_comp_for':
  65              yield last
  66          elif not last.type == 'comp_if':
  67              break
  68          last = last.children[-1]
  69  
  70  
  71  def for_stmt_defines_one_name(for_stmt):
  72      """
  73      Returns True if only one name is returned: ``for x in y``.
  74      Returns False if the for loop is more complicated: ``for x, z in y``.
  75  
  76      :returns: bool
  77      """
  78      return for_stmt.children[1].type == 'name'
  79  
  80  
  81  def get_flow_branch_keyword(flow_node, node):
  82      start_pos = node.start_pos
  83      if not (flow_node.start_pos < start_pos <= flow_node.end_pos):
  84          raise ValueError('The node is not part of the flow.')
  85  
  86      keyword = None
  87      for i, child in enumerate(flow_node.children):
  88          if start_pos < child.start_pos:
  89              return keyword
  90          first_leaf = child.get_first_leaf()
  91          if first_leaf in _FLOW_KEYWORDS:
  92              keyword = first_leaf
  93      return None
  94  
  95  
  96  def clean_scope_docstring(scope_node):
  97      """ Returns a cleaned version of the docstring token. """
  98      node = scope_node.get_doc_node()
  99      if node is not None:
 100          # TODO We have to check next leaves until there are no new
 101          # leaves anymore that might be part of the docstring. A
 102          # docstring can also look like this: ``'foo' 'bar'
 103          # Returns a literal cleaned version of the ``Token``.
 104          return cleandoc(safe_literal_eval(node.value))
 105      return ''
 106  
 107  
 108  def find_statement_documentation(tree_node):
 109      if tree_node.type == 'expr_stmt':
 110          tree_node = tree_node.parent  # simple_stmt
 111          maybe_string = tree_node.get_next_sibling()
 112          if maybe_string is not None:
 113              if maybe_string.type == 'simple_stmt':
 114                  maybe_string = maybe_string.children[0]
 115                  if maybe_string.type == 'string':
 116                      return cleandoc(safe_literal_eval(maybe_string.value))
 117      return ''
 118  
 119  
 120  def safe_literal_eval(value):
 121      first_two = value[:2].lower()
 122      if first_two[0] == 'f' or first_two in ('fr', 'rf'):
 123          # literal_eval is not able to resovle f literals. We have to do that
 124          # manually, but that's right now not implemented.
 125          return ''
 126  
 127      return literal_eval(value)
 128  
 129  
 130  def get_signature(funcdef, width=72, call_string=None,
 131                    omit_first_param=False, omit_return_annotation=False):
 132      """
 133      Generate a string signature of a function.
 134  
 135      :param width: Fold lines if a line is longer than this value.
 136      :type width: int
 137      :arg func_name: Override function name when given.
 138      :type func_name: str
 139  
 140      :rtype: str
 141      """
 142      # Lambdas have no name.
 143      if call_string is None:
 144          if funcdef.type == 'lambdef':
 145              call_string = '<lambda>'
 146          else:
 147              call_string = funcdef.name.value
 148      params = funcdef.get_params()
 149      if omit_first_param:
 150          params = params[1:]
 151      p = '(' + ''.join(param.get_code() for param in params).strip() + ')'
 152      # TODO this is pretty bad, we should probably just normalize.
 153      p = re.sub(r'\s+', ' ', p)
 154      if funcdef.annotation and not omit_return_annotation:
 155          rtype = " ->" + funcdef.annotation.get_code()
 156      else:
 157          rtype = ""
 158      code = call_string + p + rtype
 159  
 160      return '\n'.join(textwrap.wrap(code, width))
 161  
 162  
 163  def move(node, line_offset):
 164      """
 165      Move the `Node` start_pos.
 166      """
 167      try:
 168          children = node.children
 169      except AttributeError:
 170          node.line += line_offset
 171      else:
 172          for c in children:
 173              move(c, line_offset)
 174  
 175  
 176  def get_following_comment_same_line(node):
 177      """
 178      returns (as string) any comment that appears on the same line,
 179      after the node, including the #
 180      """
 181      try:
 182          if node.type == 'for_stmt':
 183              whitespace = node.children[5].get_first_leaf().prefix
 184          elif node.type == 'with_stmt':
 185              whitespace = node.children[3].get_first_leaf().prefix
 186          elif node.type == 'funcdef':
 187              # actually on the next line
 188              whitespace = node.children[4].get_first_leaf().get_next_leaf().prefix
 189          else:
 190              whitespace = node.get_last_leaf().get_next_leaf().prefix
 191      except AttributeError:
 192          return None
 193      except ValueError:
 194          # TODO in some particular cases, the tree doesn't seem to be linked
 195          # correctly
 196          return None
 197      if "#" not in whitespace:
 198          return None
 199      comment = whitespace[whitespace.index("#"):]
 200      if "\r" in comment:
 201          comment = comment[:comment.index("\r")]
 202      if "\n" in comment:
 203          comment = comment[:comment.index("\n")]
 204      return comment
 205  
 206  
 207  def is_scope(node):
 208      t = node.type
 209      if t == 'comp_for':
 210          # Starting with Python 3.8, async is outside of the statement.
 211          return node.children[1].type != 'sync_comp_for'
 212  
 213      return t in ('file_input', 'classdef', 'funcdef', 'lambdef', 'sync_comp_for')
 214  
 215  
 216  def _get_parent_scope_cache(func):
 217      cache = WeakKeyDictionary()
 218  
 219      def wrapper(parso_cache_node, node, include_flows=False):
 220          if parso_cache_node is None:
 221              return func(node, include_flows)
 222  
 223          try:
 224              for_module = cache[parso_cache_node]
 225          except KeyError:
 226              for_module = cache[parso_cache_node] = {}
 227  
 228          try:
 229              return for_module[node]
 230          except KeyError:
 231              result = for_module[node] = func(node, include_flows)
 232              return result
 233      return wrapper
 234  
 235  
 236  def get_parent_scope(node, include_flows=False):
 237      """
 238      Returns the underlying scope.
 239      """
 240      scope = node.parent
 241      if scope is None:
 242          return None  # It's a module already.
 243  
 244      while True:
 245          if is_scope(scope):
 246              if scope.type in ('classdef', 'funcdef', 'lambdef'):
 247                  index = scope.children.index(':')
 248                  if scope.children[index].start_pos >= node.start_pos:
 249                      if node.parent.type == 'param' and node.parent.name == node:
 250                          pass
 251                      elif node.parent.type == 'tfpdef' and node.parent.children[0] == node:
 252                          pass
 253                      else:
 254                          scope = scope.parent
 255                          continue
 256              return scope
 257          elif include_flows and isinstance(scope, tree.Flow):
 258              # The cursor might be on `if foo`, so the parent scope will not be
 259              # the if, but the parent of the if.
 260              if not (scope.type == 'if_stmt'
 261                      and any(n.start_pos <= node.start_pos < n.end_pos
 262                              for n in scope.get_test_nodes())):
 263                  return scope
 264  
 265          scope = scope.parent
 266  
 267  
 268  get_cached_parent_scope = _get_parent_scope_cache(get_parent_scope)
 269  
 270  
 271  def get_cached_code_lines(grammar, path):
 272      """
 273      Basically access the cached code lines in parso. This is not the nicest way
 274      to do this, but we avoid splitting all the lines again.
 275      """
 276      return get_parso_cache_node(grammar, path).lines
 277  
 278  
 279  def get_parso_cache_node(grammar, path):
 280      """
 281      This is of course not public. But as long as I control parso, this
 282      shouldn't be a problem. ~ Dave
 283  
 284      The reason for this is mostly caching. This is obviously also a sign of a
 285      broken caching architecture.
 286      """
 287      return parser_cache[grammar._hashed][path]
 288  
 289  
 290  def cut_value_at_position(leaf, position):
 291      """
 292      Cuts of the value of the leaf at position
 293      """
 294      lines = split_lines(leaf.value, keepends=True)[:position[0] - leaf.line + 1]
 295      column = position[1]
 296      if leaf.line == position[0]:
 297          column -= leaf.column
 298      if not lines:
 299          return ''
 300      lines[-1] = lines[-1][:column]
 301      return ''.join(lines)
 302  
 303  
 304  def expr_is_dotted(node):
 305      """
 306      Checks if a path looks like `name` or `name.foo.bar` and not `name()`.
 307      """
 308      if node.type == 'atom':
 309          if len(node.children) == 3 and node.children[0] == '(':
 310              return expr_is_dotted(node.children[1])
 311          return False
 312      if node.type == 'atom_expr':
 313          children = node.children
 314          if children[0] == 'await':
 315              return False
 316          if not expr_is_dotted(children[0]):
 317              return False
 318          # Check trailers
 319          return all(c.children[0] == '.' for c in children[1:])
 320      return node.type == 'name'
 321  
 322  
 323  def _function_is_x_method(decorator_checker):
 324      def wrapper(function_node):
 325          """
 326          This is a heuristic. It will not hold ALL the times, but it will be
 327          correct pretty much for anyone that doesn't try to beat it.
 328          staticmethod/classmethod are builtins and unless overwritten, this will
 329          be correct.
 330          """
 331          for decorator in function_node.get_decorators():
 332              dotted_name = decorator.children[1]
 333              if decorator_checker(dotted_name.get_code()):
 334                  return True
 335          return False
 336      return wrapper
 337  
 338  
 339  function_is_staticmethod = _function_is_x_method(lambda m: m == "staticmethod")
 340  function_is_classmethod = _function_is_x_method(lambda m: m == "classmethod")
 341  function_is_property = _function_is_x_method(
 342      lambda m: m == "property"
 343      or m == "cached_property"
 344      or (m.endswith(".setter"))
 345  )
 346