GildedRose-Refactoring-Kata/.venv/lib/python3.12/site-packages/approvaltests/inline/split_code.py
2025-06-22 13:36:01 +05:30

61 lines
2.1 KiB
Python

import re
from enum import Enum
from typing_extensions import override
class SplitCode:
def __init__(self, before_method: str, after_method: str, tab: str) -> None:
self.before_method = before_method
self.after_method = after_method
self.tab = tab
@override
def __str__(self) -> str:
return f"before:\n{self.before_method}\nafter:\n{self.after_method}\ntab: '{self.tab}'\n"
@staticmethod
def on_method(code: str, method_name: str) -> "SplitCode":
lines = code.split("\n")
before = []
after = []
tab = " "
class State(Enum):
BEFORE = 0
FIRST_LINE_OF_METHOD_BODY = 1
IN_DOCSTRING = 2
AFTER_DOCTSTRING = 3
state = State.BEFORE
for line in lines:
stripped_line = line.strip()
if stripped_line.startswith(f"def {method_name}("):
state = State.FIRST_LINE_OF_METHOD_BODY
before.append(line)
elif state == State.BEFORE:
before.append(line)
elif state == State.FIRST_LINE_OF_METHOD_BODY:
tab = re.compile(r"^\s*").match(line).group()
if stripped_line.startswith('"""'):
state = State.IN_DOCSTRING
doc_string_quotes = '"""'
elif stripped_line.startswith("'''"):
state = State.IN_DOCSTRING
doc_string_quotes = "'''"
else:
state = State.AFTER_DOCTSTRING
after.append(line)
elif state == State.IN_DOCSTRING:
if stripped_line.startswith(doc_string_quotes):
state = State.AFTER_DOCTSTRING
elif state == State.AFTER_DOCTSTRING:
after.append(line)
return SplitCode("\n".join(before), "\n".join(after), tab)
def indent(self, received_text: str) -> str:
lines = received_text.split("\n")
indented_lines = [f"{self.tab}{line}" for line in lines]
return "\n".join(indented_lines)