Add hook to remove newlines after equals signs.

Author: domdfcoding

formate/mini_hooks.py

@@ -29,17 +29,19 @@
 # stdlib
 import ast
 import re
+from collections import deque
 from typing import List
 
 # 3rd party
+import tokenize_rt  # type: ignore[import-untyped]
 from domdf_python_tools.paths import PathPlus
 from domdf_python_tools.stringlist import StringList
 from domdf_python_tools.typing import PathLike
 
 # this package
 from formate.config import wants_filename
 
-__all__ = ("noqa_reformat", "check_ast", "squish_stubs")
+__all__ = ("check_ast", "newline_after_equals", "noqa_reformat", "squish_stubs")
 
 
 def noqa_reformat(source: str) -> str:
@@ -246,3 +248,40 @@ def _reformat_blocks(blocks: List[List[str]]) -> StringList:
 	output.blankline(ensure_single=True)
 
 	return output
+
+
+def newline_after_equals(source: str) -> str:
+	"""
+	Removes newlines immediately after equals signs.
+
+	.. versionadded:: 1.2.0
+
+	:param source: The source to check.
+
+	:return: The reformatted source.
+	"""
+
+	original_tokens = deque(tokenize_rt.src_to_tokens(source))
+	tokens = []
+
+	while original_tokens:
+		token = original_tokens.popleft()
+
+		if token.name == "OP" and token.src == '=':
+			while True:
+
+				# TODO: handle spaces around equals (e.g. in function signature)
+				# Look ahead until there's a non-whitespace token.
+				next_token = original_tokens.popleft()
+				if next_token.name in {"UNIMPORTANT_WS", "NL"}:
+					pass
+				else:
+					break
+
+			tokens.append(token)
+			tokens.append(next_token)
+
+		else:
+			tokens.append(token)
+
+	return tokenize_rt.tokens_to_src(tokens)

pyproject.toml

@@ -31,6 +31,7 @@ reformat-generics = "formate.reformat_generics:reformat_generics"
 dynamic_quotes = "formate.dynamic_quotes:dynamic_quotes"
 noqa_reformat = "formate.mini_hooks:noqa_reformat"
 squish_stubs = "formate.mini_hooks:squish_stubs"
+newline_after_equals = "formate.mini_hooks:newline_after_equals"
 ellipsis_reformat = "formate.ellipses:ellipsis_reformat"
 collections-import-rewrite = "formate.imports:rewrite_collections_abc_imports"
 isort = "formate:isort_hook"

repo_helper.yml

@@ -68,6 +68,7 @@ entry_points:
    - "dynamic_quotes=formate.dynamic_quotes:dynamic_quotes"
    - "noqa_reformat=formate.mini_hooks:noqa_reformat"
    - "squish_stubs=formate.mini_hooks:squish_stubs"
+   - "newline_after_equals=formate.mini_hooks:newline_after_equals"
    - "ellipsis_reformat=formate.ellipses:ellipsis_reformat"
    - "collections-import-rewrite=formate.imports:rewrite_collections_abc_imports"
    - "isort=formate:isort_hook"

requirements.txt

@@ -8,5 +8,6 @@ dom-toml>=2.2.0
 domdf-python-tools>=2.5.0
 isort>=5.5.2
 prettyprinter>=0.18.0
+tokenize-rt>=5.0.0
 typing-extensions>=3.7.4.3
 yapf<0.43.0,>=0.30.0

tests/test_mini_hooks.py

@@ -6,7 +6,7 @@
 from coincidence.regressions import AdvancedFileRegressionFixture
 
 # this package
-from formate.mini_hooks import check_ast, noqa_reformat, squish_stubs
+from formate.mini_hooks import check_ast, newline_after_equals, noqa_reformat, squish_stubs
 
 
 def test_noqa_reformat():
@@ -254,3 +254,33 @@ def foo():
 
 	new_code = squish_stubs(code, formate_filename="code.py")
 	assert new_code == code
+
+
+newline_after_equals_src = """
+def foo():
+
+	with pytest.raises(
+			ValueError,
+			match=
+			"This is a really long error message that exceeds the line length limit so gets pushed down!!!!!!",
+			):
+		pass
+
+"""
+
+no_newline_after_equals_src = """
+def foo():
+
+	with pytest.raises(
+			ValueError,
+			match="This is a shorter error message that's still too long to fit on one line.",
+			):
+		pass
+
+"""
+
+
+def test_newline_after_equals(advanced_file_regression: AdvancedFileRegressionFixture):
+
+	advanced_file_regression.check(newline_after_equals(newline_after_equals_src))
+	assert newline_after_equals(no_newline_after_equals_src) == no_newline_after_equals_src

tests/test_mini_hooks_/test_newline_after_equals.txt

@@ -0,0 +1,8 @@
+
+def foo():
+
+	with pytest.raises(
+			ValueError,
+			match="This is a really long error message that exceeds the line length limit so gets pushed down!!!!!!",
+			):
+		pass