<commit_before># Copyright 2022 <NAME>, All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import ast

from yapf.pyparser import pyparser_utils as pyutils
from yapf.yapflib import split_penalty
from yapf.yapflib import style
from yapf.yapflib import subtypes


# This is a skeleton of an AST visitor.
class SplitPenalty(ast.NodeVisitor):
  """Compute split penalties between tokens."""

  def __init__(self, logical_lines):
    super(SplitPenalty, self).__init__()
    self.logical_lines = logical_lines

    # We never want to split before a colon or comma.
    for logical_line in logical_lines:
      for token in logical_line.tokens:
        if token.value in frozenset({',', ':'}):
          token.split_penalty = split_penalty.UNBREAKABLE

  def _GetTokens(self, node):
    return pyutils.GetTokens(self.logical_lines, node)

  ############################################################################
  # Statements                                                               #
  ############################################################################

  def visit_FunctionDef(self, node):
    # FunctionDef(name=Name,
    #             args=arguments(
    #                 posonlyargs=[],
    #                 args=[],
    #                 vararg=[],
    #                 kwonlyargs=[],
    #                 kw_defaults=[],
    #                 defaults=[]),
    #             body=[...],
    #             decorator_list=[Call_1, Call_2, ..., Call_n],
    #             keywords=[])
    tokens = self._GetTokens(node)
    for decorator in node.decorator_list:
      # The decorator token list begins after the '@'. The body of the decorator
      # is formatted like a normal "call."
      subrange = pyutils.GetTokensInSubRange(tokens, decorator)
      # Don't split after the '@'.
      subrange[0].split_penalty = split_penalty.UNBREAKABLE

    if node.returns:
      start_index = pyutils.GetTokenIndex(tokens,
                                          pyutils.TokenStart(node.returns))
      _IncreasePenalty(tokens[start_index - 1:start_index + 1],
                       split_penalty.VERY_STRONGLY_CONNECTED)
      end_index = pyutils.GetTokenIndex(tokens, pyutils.TokenEnd(node.returns))
      _IncreasePenalty(tokens[start_index + 1:end_index],
                       split_penalty.STRONGLY_CONNECTED)

    return self.generic_visit(node)

  def visit_AsyncFunctionDef(self, node):
    # AsyncFunctionDef(name=Name,
    #                  args=arguments(
    #                      posonlyargs=[],
    #                      args=[],
    #                      vararg=[],
    #                      kwonlyargs=[],
    #                      kw_defaults=[],
    #                      defaults=[]),
    #                  body=[...],
    #                  decorator_list=[Expr_1, Expr_2, ..., Expr_n],
    #                  keywords=[])
    tokens = self._GetTokens(node)
    _IncreasePenalty(tokens[1], split_penalty.UNBREAKABLE)

    return self.visit_FunctionDef(node)

  def visit_ClassDef(self, node):
    # ClassDef(name=Name,
    #          bases=[Expr_1, Expr_2, ..., Expr_n],
    #          keywords=[],
    #          body=[],
    #          decorator_list=[Expr_1, Expr_2, ..., Expr_m])
    for base in node.bases:
      tokens = self._GetTokens(base)
      _IncreasePenalty(tokens[1:], split_penalty.EXPR)

    for decorator in node.decorator_list:
      # Don't split after the '@'.
      tokens = self._GetTokens(decorator)
      tokens[0].split_penalty = split_penalty.UNBREAKABLE

    return self.generic_visit(node)

  def visit_Return(self, node):
    # Return(value=Expr)
    tokens = self._GetTokens(node)
    _IncreasePenalty(tokens[1:], split_penalty.EXPR)

    return self.generic_visit(node)

  def visit_Delete(self, node):
    # Delete(targets=[Expr_1, Expr_2, ..., Expr_n])
    for target in node.targets:
      tokens = self._GetTokens(target)
      _IncreasePenalty(tokens[1:], split_penalty.EXPR)

    return self.generic_visit(node)

  def visit_Assign(self, node):
    # Assign(targets=[Expr_1, Expr_2, ..., Expr_n],
    #        value=Expr)<commit_msg>Decorators aren't in the same range as the function def node<commit_after># Copyright 2022 <NAME>, All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import ast

from yapf.pyparser import pyparser_utils as pyutils
from yapf.yapflib import split_penalty
from yapf.yapflib import style
from yapf.yapflib import subtypes


# This is a skeleton of an AST visitor.
class SplitPenalty(ast.NodeVisitor):
  """Compute split penalties between tokens."""

  def __init__(self, logical_lines):
    super(SplitPenalty, self).__init__()
    self.logical_lines = logical_lines

    # We never want to split before a colon or comma.
    for logical_line in logical_lines:
      for token in logical_line.tokens:
        if token.value in frozenset({',', ':'}):
          token.split_penalty = split_penalty.UNBREAKABLE

  def _GetTokens(self, node):
    return pyutils.GetTokens(self.logical_lines, node)

  ############################################################################
  # Statements                                                               #
  ############################################################################

  def visit_FunctionDef(self, node):
    # FunctionDef(name=Name,
    #             args=arguments(
    #                 posonlyargs=[],
    #                 args=[],
    #                 vararg=[],
    #                 kwonlyargs=[],
    #                 kw_defaults=[],
    #                 defaults=[]),
    #             body=[...],
    #             decorator_list=[Call_1, Call_2, ..., Call_n],
    #             keywords=[])
    tokens = self._GetTokens(node)
    for decorator in node.decorator_list:
      # The decorator token list begins after the '@'. The body of the decorator
      # is formatted like a normal "call."
      decorator_range = self._GetTokens(decorator)
      # Don't split after the '@'.
      decorator_range[0].split_penalty = split_penalty.UNBREAKABLE

    if node.returns:
      start_index = pyutils.GetTokenIndex(tokens,
                                          pyutils.TokenStart(node.returns))
      _IncreasePenalty(tokens[start_index - 1:start_index + 1],
                       split_penalty.VERY_STRONGLY_CONNECTED)
      end_index = pyutils.GetTokenIndex(tokens, pyutils.TokenEnd(node.returns))
      _IncreasePenalty(tokens[start_index + 1:end_index],
                       split_penalty.STRONGLY_CONNECTED)

    return self.generic_visit(node)

  def visit_AsyncFunctionDef(self, node):
    # AsyncFunctionDef(name=Name,
    #                  args=arguments(
    #                      posonlyargs=[],
    #                      args=[],
    #                      vararg=[],
    #                      kwonlyargs=[],
    #                      kw_defaults=[],
    #                      defaults=[]),
    #                  body=[...],
    #                  decorator_list=[Expr_1, Expr_2, ..., Expr_n],
    #                  keywords=[])
    tokens = self._GetTokens(node)
    _IncreasePenalty(tokens[1], split_penalty.UNBREAKABLE)

    return self.visit_FunctionDef(node)

  def visit_ClassDef(self, node):
    # ClassDef(name=Name,
    #          bases=[Expr_1, Expr_2, ..., Expr_n],
    #          keywords=[],
    #          body=[],
    #          decorator_list=[Expr_1, Expr_2, ..., Expr_m])
    for base in node.bases:
      tokens = self._GetTokens(base)
      _IncreasePenalty(tokens[1:], split_penalty.EXPR)

    for decorator in node.decorator_list:
      # Don't split after the '@'.
      decorator_range = self._GetTokens(decorator)
      decorator_range[0].split_penalty = split_penalty.UNBREAKABLE

    return self.generic_visit(node)

  def visit_Return(self, node):
    # Return(value=Expr)
    tokens = self._GetTokens(node)
    _IncreasePenalty(tokens[1:], split_penalty.EXPR)

    return self.generic_visit(node)

  def visit_Delete(self, node):
    # Delete(targets=[Expr_1, Expr_2, ..., Expr_n])
    for target in node.targets:
      tokens = self._GetTokens(target)
      _IncreasePenalty(tokens[1:], split_penalty.EXPR)

    return self.generic_visit(node)

  def visit_Assign(self, node):
    # Assign(targets=[Expr_1, Expr_2, ..., Expr_n],
    #        value=Expr)
{'fraction_non_alphanumeric': 0.13261917514729513, 'fraction_numerical': 0.006213176218532405, 'mean_word_length': 2.3558590941768514, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2963658', 'n_tokens_mistral': 2817, 'n_tokens_neox': 2660, 'n_words': 634}