Coverage for src\sql_organizer\file_formatter.py: 100%

1import re 

2from abc import ABC, abstractmethod 

3from dataclasses import dataclass 

4 

5from sql_organizer.file import SqlFile 

6 

7COMMENT_REGEX = re.compile(r"--\s*use", flags=re.I + re.M) 

8 

9 

10class SqlFileFormatter(ABC): 

11 @abstractmethod 

12 def format_file(self, sql_file: SqlFile) -> str: ... 

13 

14 def and_(self, formatter: "SqlFileFormatter") -> "FormatterCombination": 

15 return FormatterCombination(formatters=(self, formatter)) 

16 

17 

18@dataclass(slots=True) 

19class FormatterCombination(SqlFileFormatter): 

20 formatters: tuple[SqlFileFormatter, ...] 

21 

22 def format_file(self, sql_file: SqlFile) -> str: 

23 return "\n".join(formater.format_file(sql_file) for formater in self.formatters) 

24 

25 def and_(self, formatter: SqlFileFormatter) -> "FormatterCombination": 

26 return FormatterCombination(formatters=(*self.formatters, formatter)) 

27 

28 

29class NameFormatter(SqlFileFormatter): 

30 def format_file(self, sql_file: SqlFile) -> str: 

31 return f"-- {sql_file.file_name}" 

32 

33 

34class SkipLine(SqlFileFormatter): 

35 def format_file(self, sql_file: SqlFile) -> str: 

36 return "\n" 

37 

38 

39class PlainSqlText(SqlFileFormatter): 

40 def format_file(self, sql_file: SqlFile) -> str: 

41 return sql_file.sql_text 

42 

43 

44class BreakLine(SqlFileFormatter): 

45 def format_file(self, sql_file: SqlFile) -> str: 

46 return "--" + "_" * 20 + f" End of {sql_file.file_name} " + "_" * 20 + "--" 

47 

48 

49@dataclass(slots=True) 

50class UseCommentRemover(SqlFileFormatter): 

51 formatter: SqlFileFormatter 

52 

53 def format_file(self, sql_file: SqlFile) -> str: 

54 return COMMENT_REGEX.sub("USE", self.formatter.format_file(sql_file)) 

55 

56 

57def get_standard_formatter(remove_comments: bool) -> SqlFileFormatter: 

58 return ( 

59 NameFormatter() 

60 .and_(SkipLine()) 

61 .and_(SkipLine()) 

62 .and_( 

63 UseCommentRemover(formatter=PlainSqlText()) 

64 if remove_comments 

65 else PlainSqlText() 

66 ) 

67 .and_(SkipLine()) 

68 .and_(BreakLine()) 

69 .and_(SkipLine()) 

70 .and_(SkipLine()) 

71 )