sqlnbfmt 0.1.8__tar.gz → 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {sqlnbfmt-0.1.8/sqlnbfmt.egg-info → sqlnbfmt-0.2.0}/PKG-INFO +8 -8
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/README.md +4 -4
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/pyproject.toml +9 -6
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/sqlnbfmt/formatter.py +140 -154
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0/sqlnbfmt.egg-info}/PKG-INFO +8 -8
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/sqlnbfmt.egg-info/SOURCES.txt +1 -0
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/sqlnbfmt.egg-info/requires.txt +0 -2
- sqlnbfmt-0.2.0/tests/test_eval.py +81 -0
- sqlnbfmt-0.2.0/tests/test_formatter.py +237 -0
- sqlnbfmt-0.1.8/tests/test_formatter.py +0 -201
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/LICENSE +0 -0
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/setup.cfg +0 -0
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/sqlnbfmt/__init__.py +0 -0
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/sqlnbfmt.egg-info/dependency_links.txt +0 -0
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/sqlnbfmt.egg-info/entry_points.txt +0 -0
- {sqlnbfmt-0.1.8 → sqlnbfmt-0.2.0}/sqlnbfmt.egg-info/top_level.txt +0 -0
|
@@ -1,19 +1,19 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: sqlnbfmt
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
Summary: A tool to format SQL code in Jupyter notebooks.
|
|
5
5
|
Author-email: Qing <qingye779@gmail.com>
|
|
6
6
|
License: MIT
|
|
7
7
|
Classifier: Programming Language :: Python :: 3
|
|
8
8
|
Classifier: Programming Language :: Python :: 3.12
|
|
9
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
9
10
|
Requires-Python: >=3.12
|
|
10
11
|
Description-Content-Type: text/markdown
|
|
11
12
|
License-File: LICENSE
|
|
12
|
-
Requires-Dist: astor>=0.8.1
|
|
13
|
-
Requires-Dist: black>=24.10.0
|
|
14
13
|
Requires-Dist: nbformat>=5.10.4
|
|
15
14
|
Requires-Dist: pyyaml>=6.0.2
|
|
16
15
|
Requires-Dist: sqlglot>=25.27.0
|
|
16
|
+
Dynamic: license-file
|
|
17
17
|
|
|
18
18
|
# sqlnbfmt
|
|
19
19
|
|
|
@@ -27,10 +27,10 @@ A SQL formatter designed specifically for Jupyter Notebooks. `sqlnbfmt` automati
|
|
|
27
27
|
|
|
28
28
|
- 🎯 **Smart SQL Detection**: Automatically identifies and formats SQL queries in code cells and magic SQL cells
|
|
29
29
|
- 🌳 **AST-Powered**: Uses Abstract Syntax Tree parsing for accurate SQL string identification
|
|
30
|
-
- 🔒 **Safe Formatting**: Preserves query parameters (e.g., `%s`, `?`) and SQL comments
|
|
30
|
+
- 🔒 **Safe Formatting**: Preserves Python comments, query parameters (e.g., `%s`, `?`), and SQL comments
|
|
31
31
|
- ⚙️ **Highly Configurable**: Customize formatting through YAML configuration
|
|
32
32
|
- 🔄 **Pre-commit Ready**: Seamlessly integrates with pre-commit hooks
|
|
33
|
-
- 📦 **
|
|
33
|
+
- 📦 **Lightweight**: Only three runtime dependencies (sqlglot, nbformat, pyyaml)
|
|
34
34
|
|
|
35
35
|
## Installation
|
|
36
36
|
|
|
@@ -63,14 +63,14 @@ pip install pre-commit
|
|
|
63
63
|
```yaml
|
|
64
64
|
repos:
|
|
65
65
|
- repo: https://github.com/flyersworder/sqlnbfmt
|
|
66
|
-
rev: v0.
|
|
66
|
+
rev: v0.2.0
|
|
67
67
|
hooks:
|
|
68
68
|
- id: sqlnbfmt
|
|
69
69
|
name: sqlnbfmt
|
|
70
70
|
types: [jupyter]
|
|
71
71
|
args: [--config, config.yaml, --dialect, postgres]
|
|
72
72
|
```
|
|
73
|
-
Please
|
|
73
|
+
Please run the following command for help:
|
|
74
74
|
|
|
75
75
|
```bash
|
|
76
76
|
sqlnbfmt -h
|
|
@@ -10,10 +10,10 @@ A SQL formatter designed specifically for Jupyter Notebooks. `sqlnbfmt` automati
|
|
|
10
10
|
|
|
11
11
|
- 🎯 **Smart SQL Detection**: Automatically identifies and formats SQL queries in code cells and magic SQL cells
|
|
12
12
|
- 🌳 **AST-Powered**: Uses Abstract Syntax Tree parsing for accurate SQL string identification
|
|
13
|
-
- 🔒 **Safe Formatting**: Preserves query parameters (e.g., `%s`, `?`) and SQL comments
|
|
13
|
+
- 🔒 **Safe Formatting**: Preserves Python comments, query parameters (e.g., `%s`, `?`), and SQL comments
|
|
14
14
|
- ⚙️ **Highly Configurable**: Customize formatting through YAML configuration
|
|
15
15
|
- 🔄 **Pre-commit Ready**: Seamlessly integrates with pre-commit hooks
|
|
16
|
-
- 📦 **
|
|
16
|
+
- 📦 **Lightweight**: Only three runtime dependencies (sqlglot, nbformat, pyyaml)
|
|
17
17
|
|
|
18
18
|
## Installation
|
|
19
19
|
|
|
@@ -46,14 +46,14 @@ pip install pre-commit
|
|
|
46
46
|
```yaml
|
|
47
47
|
repos:
|
|
48
48
|
- repo: https://github.com/flyersworder/sqlnbfmt
|
|
49
|
-
rev: v0.
|
|
49
|
+
rev: v0.2.0
|
|
50
50
|
hooks:
|
|
51
51
|
- id: sqlnbfmt
|
|
52
52
|
name: sqlnbfmt
|
|
53
53
|
types: [jupyter]
|
|
54
54
|
args: [--config, config.yaml, --dialect, postgres]
|
|
55
55
|
```
|
|
56
|
-
Please
|
|
56
|
+
Please run the following command for help:
|
|
57
57
|
|
|
58
58
|
```bash
|
|
59
59
|
sqlnbfmt -h
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "sqlnbfmt"
|
|
3
|
-
version = "0.
|
|
3
|
+
version = "0.2.0"
|
|
4
4
|
description = "A tool to format SQL code in Jupyter notebooks."
|
|
5
5
|
authors = [
|
|
6
6
|
{ name = "Qing", email = "qingye779@gmail.com" }
|
|
@@ -10,19 +10,19 @@ readme = "README.md"
|
|
|
10
10
|
requires-python = ">=3.12"
|
|
11
11
|
classifiers = [
|
|
12
12
|
"Programming Language :: Python :: 3",
|
|
13
|
-
"Programming Language :: Python :: 3.12"
|
|
13
|
+
"Programming Language :: Python :: 3.12",
|
|
14
|
+
"Programming Language :: Python :: 3.13"
|
|
14
15
|
]
|
|
15
16
|
dependencies = [
|
|
16
|
-
"astor>=0.8.1",
|
|
17
|
-
"black>=24.10.0",
|
|
18
17
|
"nbformat>=5.10.4",
|
|
19
18
|
"pyyaml>=6.0.2",
|
|
20
19
|
"sqlglot>=25.27.0",
|
|
21
20
|
]
|
|
22
21
|
|
|
23
|
-
[
|
|
24
|
-
dev
|
|
22
|
+
[dependency-groups]
|
|
23
|
+
dev = [
|
|
25
24
|
"pytest-cov>=5.0.0",
|
|
25
|
+
"pytest-timeout>=2.3.1",
|
|
26
26
|
"pytest>=8.3.3",
|
|
27
27
|
]
|
|
28
28
|
|
|
@@ -38,5 +38,8 @@ markers = [
|
|
|
38
38
|
"integration: mark a test as an integration test."
|
|
39
39
|
]
|
|
40
40
|
|
|
41
|
+
[tool.ruff]
|
|
42
|
+
extend-exclude = ["tests/eval/*/"]
|
|
43
|
+
|
|
41
44
|
[tool.setuptools.package-data]
|
|
42
45
|
sqlnbfmt = [".pre-commit-hooks.yaml"]
|
|
@@ -6,13 +6,20 @@ from dataclasses import dataclass
|
|
|
6
6
|
from pathlib import Path
|
|
7
7
|
from typing import Set, Optional, Dict, Union, Tuple
|
|
8
8
|
|
|
9
|
-
import astor
|
|
10
|
-
import black
|
|
11
9
|
import nbformat
|
|
12
10
|
import yaml
|
|
13
11
|
from sqlglot import parse_one, errors
|
|
14
12
|
|
|
15
13
|
|
|
14
|
+
@dataclass
|
|
15
|
+
class SQLReplacement:
|
|
16
|
+
"""A text replacement to apply to the source code."""
|
|
17
|
+
|
|
18
|
+
start_offset: int
|
|
19
|
+
end_offset: int
|
|
20
|
+
new_text: str
|
|
21
|
+
|
|
22
|
+
|
|
16
23
|
@dataclass
|
|
17
24
|
class FormattingConfig:
|
|
18
25
|
"""Configuration for SQL formatting."""
|
|
@@ -21,7 +28,6 @@ class FormattingConfig:
|
|
|
21
28
|
function_names: Set[str]
|
|
22
29
|
sql_decorators: Set[str]
|
|
23
30
|
single_line_threshold: int = 80
|
|
24
|
-
preserve_comments: bool = True
|
|
25
31
|
indent_width: int = 4
|
|
26
32
|
|
|
27
33
|
|
|
@@ -33,7 +39,7 @@ class SQLFormattingError(Exception):
|
|
|
33
39
|
|
|
34
40
|
def setup_logging(level: str = "INFO") -> logging.Logger:
|
|
35
41
|
"""Sets up logging with the specified level."""
|
|
36
|
-
logger = logging.getLogger(
|
|
42
|
+
logger = logging.getLogger("formatter")
|
|
37
43
|
logger.setLevel(level.upper())
|
|
38
44
|
if not logger.handlers:
|
|
39
45
|
handler = logging.StreamHandler(sys.stdout)
|
|
@@ -59,9 +65,6 @@ def load_config(config_path: Union[str, Path] = "config.yaml") -> FormattingConf
|
|
|
59
65
|
single_line_threshold=config.get("formatting_options", {}).get(
|
|
60
66
|
"single_line_threshold", 80
|
|
61
67
|
),
|
|
62
|
-
preserve_comments=config.get("formatting_options", {}).get(
|
|
63
|
-
"preserve_comments", True
|
|
64
|
-
),
|
|
65
68
|
indent_width=config.get("formatting_options", {}).get(
|
|
66
69
|
"indent_width", 4
|
|
67
70
|
),
|
|
@@ -70,6 +73,22 @@ def load_config(config_path: Union[str, Path] = "config.yaml") -> FormattingConf
|
|
|
70
73
|
raise yaml.YAMLError(f"Error parsing the configuration file: {e}")
|
|
71
74
|
|
|
72
75
|
|
|
76
|
+
def build_line_offsets(source: str) -> list:
|
|
77
|
+
"""Returns list where offsets[i] is the character offset of line i+1 in source."""
|
|
78
|
+
offsets = [0]
|
|
79
|
+
for i, ch in enumerate(source):
|
|
80
|
+
if ch == "\n":
|
|
81
|
+
offsets.append(i + 1)
|
|
82
|
+
return offsets
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def node_offsets(node: ast.AST, line_offsets: list) -> Tuple[int, int]:
|
|
86
|
+
"""Converts AST node line/col positions to (start, end) character offsets."""
|
|
87
|
+
start = line_offsets[node.lineno - 1] + node.col_offset
|
|
88
|
+
end = line_offsets[node.end_lineno - 1] + node.end_col_offset
|
|
89
|
+
return start, end
|
|
90
|
+
|
|
91
|
+
|
|
73
92
|
def format_sql_code(
|
|
74
93
|
sql_code: str,
|
|
75
94
|
dialect: Optional[str],
|
|
@@ -95,7 +114,7 @@ def format_sql_code(
|
|
|
95
114
|
str: The formatted SQL code.
|
|
96
115
|
"""
|
|
97
116
|
try:
|
|
98
|
-
logger = logging.getLogger(
|
|
117
|
+
logger = logging.getLogger("formatter")
|
|
99
118
|
|
|
100
119
|
if not sql_code.strip():
|
|
101
120
|
return sql_code
|
|
@@ -112,7 +131,7 @@ def format_sql_code(
|
|
|
112
131
|
placeholder_mapping[temp_placeholder] = placeholder
|
|
113
132
|
|
|
114
133
|
# Handle automatic placeholders (%s, ?)
|
|
115
|
-
auto_placeholder_pattern = re.compile(r
|
|
134
|
+
auto_placeholder_pattern = re.compile(r"%s|\?")
|
|
116
135
|
auto_placeholders = auto_placeholder_pattern.findall(temp_sql)
|
|
117
136
|
auto_placeholder_mapping = {}
|
|
118
137
|
for idx, ph in enumerate(auto_placeholders):
|
|
@@ -125,9 +144,7 @@ def format_sql_code(
|
|
|
125
144
|
# Parse and format SQL
|
|
126
145
|
parsed = parse_one(temp_sql, read=dialect)
|
|
127
146
|
formatted_sql = parsed.sql(
|
|
128
|
-
pretty=not force_single_line,
|
|
129
|
-
indent=config.indent_width,
|
|
130
|
-
dialect=dialect
|
|
147
|
+
pretty=not force_single_line, indent=config.indent_width, dialect=dialect
|
|
131
148
|
)
|
|
132
149
|
|
|
133
150
|
# Apply formatting based on context
|
|
@@ -143,14 +160,22 @@ def format_sql_code(
|
|
|
143
160
|
if placeholders:
|
|
144
161
|
for temp_placeholder, original_placeholder in placeholder_mapping.items():
|
|
145
162
|
# Remove quotes around placeholders if any
|
|
146
|
-
formatted_sql = formatted_sql.replace(
|
|
147
|
-
|
|
163
|
+
formatted_sql = formatted_sql.replace(
|
|
164
|
+
f"'{temp_placeholder}'", temp_placeholder
|
|
165
|
+
)
|
|
166
|
+
formatted_sql = formatted_sql.replace(
|
|
167
|
+
f'"{temp_placeholder}"', temp_placeholder
|
|
168
|
+
)
|
|
148
169
|
# Replace temp placeholders with original placeholders
|
|
149
|
-
formatted_sql = formatted_sql.replace(
|
|
170
|
+
formatted_sql = formatted_sql.replace(
|
|
171
|
+
temp_placeholder, original_placeholder
|
|
172
|
+
)
|
|
150
173
|
|
|
151
174
|
# Restore automatic placeholders
|
|
152
175
|
for temp_placeholder, original_placeholder in auto_placeholder_mapping.items():
|
|
153
|
-
formatted_sql = formatted_sql.replace(
|
|
176
|
+
formatted_sql = formatted_sql.replace(
|
|
177
|
+
temp_placeholder, original_placeholder
|
|
178
|
+
)
|
|
154
179
|
|
|
155
180
|
# Logging for debugging purposes
|
|
156
181
|
logger.debug(f"Original SQL:\n{sql_code}")
|
|
@@ -164,16 +189,22 @@ def format_sql_code(
|
|
|
164
189
|
raise SQLFormattingError(f"Unexpected error during SQL formatting: {e}")
|
|
165
190
|
|
|
166
191
|
|
|
167
|
-
class SQLStringFormatter(ast.
|
|
168
|
-
"""AST
|
|
192
|
+
class SQLStringFormatter(ast.NodeVisitor):
|
|
193
|
+
"""AST NodeVisitor that collects SQL string replacements."""
|
|
169
194
|
|
|
170
195
|
def __init__(
|
|
171
|
-
self,
|
|
196
|
+
self,
|
|
197
|
+
config: FormattingConfig,
|
|
198
|
+
dialect: Optional[str],
|
|
199
|
+
logger: logging.Logger,
|
|
200
|
+
line_offsets: list,
|
|
172
201
|
):
|
|
173
202
|
super().__init__()
|
|
174
203
|
self.config = config
|
|
175
204
|
self.dialect = dialect
|
|
176
205
|
self.logger = logger
|
|
206
|
+
self.line_offsets = line_offsets
|
|
207
|
+
self.replacements: list = []
|
|
177
208
|
self.changed = False
|
|
178
209
|
|
|
179
210
|
def is_likely_sql(self, code: str) -> bool:
|
|
@@ -215,7 +246,7 @@ class SQLStringFormatter(ast.NodeTransformer):
|
|
|
215
246
|
if isinstance(value, ast.Constant):
|
|
216
247
|
sql_parts.append(value.value)
|
|
217
248
|
elif isinstance(value, ast.FormattedValue):
|
|
218
|
-
expr =
|
|
249
|
+
expr = ast.unparse(value.value)
|
|
219
250
|
placeholder = f"PLACEHOLDER_{placeholder_counter}"
|
|
220
251
|
sql_parts.append(placeholder)
|
|
221
252
|
placeholders[placeholder] = expr
|
|
@@ -224,128 +255,52 @@ class SQLStringFormatter(ast.NodeTransformer):
|
|
|
224
255
|
return "".join(sql_parts), placeholders
|
|
225
256
|
|
|
226
257
|
def format_sql_node(
|
|
227
|
-
self,
|
|
228
|
-
node: Union[ast.Constant, ast.JoinedStr],
|
|
258
|
+
self,
|
|
259
|
+
node: Union[ast.Constant, ast.JoinedStr],
|
|
229
260
|
force_single_line: bool = False,
|
|
230
|
-
in_function: bool = False
|
|
231
|
-
) -> Optional[
|
|
232
|
-
"""Formats SQL code in AST nodes."""
|
|
261
|
+
in_function: bool = False,
|
|
262
|
+
) -> Optional[str]:
|
|
263
|
+
"""Formats SQL code in AST nodes. Returns replacement text or None."""
|
|
233
264
|
try:
|
|
265
|
+
is_fstring = isinstance(node, ast.JoinedStr)
|
|
266
|
+
placeholders = {}
|
|
267
|
+
|
|
234
268
|
if isinstance(node, ast.Constant) and isinstance(node.value, str):
|
|
235
269
|
if not self.is_likely_sql(node.value):
|
|
236
270
|
return None
|
|
271
|
+
sql_str = node.value
|
|
272
|
+
elif isinstance(node, ast.JoinedStr):
|
|
273
|
+
sql_str, placeholders = self.extract_fstring_parts(node)
|
|
274
|
+
if not sql_str or not self.is_likely_sql(sql_str):
|
|
275
|
+
return None
|
|
276
|
+
else:
|
|
277
|
+
return None
|
|
237
278
|
|
|
279
|
+
try:
|
|
238
280
|
formatted_sql = format_sql_code(
|
|
239
|
-
|
|
281
|
+
sql_str,
|
|
240
282
|
self.dialect,
|
|
241
283
|
self.config,
|
|
284
|
+
placeholders=placeholders if placeholders else None,
|
|
242
285
|
force_single_line=force_single_line,
|
|
243
286
|
)
|
|
287
|
+
except SQLFormattingError:
|
|
288
|
+
return None
|
|
244
289
|
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
if in_function and '\n' in formatted_sql:
|
|
248
|
-
# Apply Black-style indentation for function arguments
|
|
249
|
-
lines = formatted_sql.split('\n')
|
|
250
|
-
# Indent all lines including the first SQL line
|
|
251
|
-
indented_lines = [
|
|
252
|
-
" " * 4 + line if line.strip() else line
|
|
253
|
-
for line in lines
|
|
254
|
-
]
|
|
255
|
-
formatted_sql = '\n'.join(indented_lines)
|
|
256
|
-
|
|
257
|
-
# Reconstruct the string with appropriate quotes
|
|
258
|
-
formatted_str = f'"""\n{formatted_sql}\n{" " * 4}"""'
|
|
259
|
-
else:
|
|
260
|
-
# Ensure newlines are added around the formatted SQL
|
|
261
|
-
formatted_str = f'"""\n{formatted_sql}\n"""'
|
|
262
|
-
|
|
263
|
-
formatted_node = ast.parse(formatted_str).body[0].value
|
|
264
|
-
return formatted_node
|
|
290
|
+
if formatted_sql == sql_str:
|
|
291
|
+
return None
|
|
265
292
|
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
|
|
272
|
-
formatted_sql =
|
|
273
|
-
sql_str,
|
|
274
|
-
self.dialect,
|
|
275
|
-
self.config,
|
|
276
|
-
placeholders=placeholders,
|
|
277
|
-
force_single_line=force_single_line,
|
|
278
|
-
)
|
|
279
|
-
except SQLFormattingError:
|
|
280
|
-
return None
|
|
293
|
+
# Restore placeholders as f-string expressions
|
|
294
|
+
if placeholders:
|
|
295
|
+
for ph in placeholders:
|
|
296
|
+
formatted_sql = formatted_sql.replace(f"'{ph}'", ph)
|
|
297
|
+
formatted_sql = formatted_sql.replace(f'"{ph}"', ph)
|
|
298
|
+
for ph, expr in placeholders.items():
|
|
299
|
+
formatted_sql = formatted_sql.replace(ph, f"{{{expr}}}")
|
|
281
300
|
|
|
282
|
-
|
|
283
|
-
|
|
284
|
-
if not placeholder:
|
|
285
|
-
continue
|
|
286
|
-
formatted_sql = formatted_sql.replace(f"'{placeholder}'", placeholder)
|
|
287
|
-
formatted_sql = formatted_sql.replace(f'"{placeholder}"', placeholder)
|
|
288
|
-
|
|
289
|
-
if formatted_sql != sql_str:
|
|
290
|
-
self.changed = True
|
|
291
|
-
# Apply Black-style indentation for function arguments
|
|
292
|
-
if in_function and '\n' in formatted_sql:
|
|
293
|
-
lines = formatted_sql.split('\n')
|
|
294
|
-
# Indent all lines including the first SQL line
|
|
295
|
-
indented_lines = [
|
|
296
|
-
" " * 4 + line if line.strip() else line
|
|
297
|
-
for line in lines
|
|
298
|
-
]
|
|
299
|
-
formatted_sql = '\n'.join(indented_lines)
|
|
300
|
-
|
|
301
|
-
# Reconstruct the f-string
|
|
302
|
-
new_values = []
|
|
303
|
-
if not placeholders:
|
|
304
|
-
new_values.append(ast.Constant(value=formatted_sql))
|
|
305
|
-
else:
|
|
306
|
-
pattern = re.compile('|'.join(re.escape(k) for k in placeholders.keys() if k))
|
|
307
|
-
idx = 0
|
|
308
|
-
while idx < len(formatted_sql):
|
|
309
|
-
match = pattern.search(formatted_sql, idx)
|
|
310
|
-
if match:
|
|
311
|
-
if match.start() > idx:
|
|
312
|
-
new_values.append(ast.Constant(value=formatted_sql[idx:match.start()]))
|
|
313
|
-
placeholder = match.group()
|
|
314
|
-
expr_str = placeholders[placeholder].strip()
|
|
315
|
-
try:
|
|
316
|
-
expr_ast = ast.parse(expr_str, mode='eval').body
|
|
317
|
-
except SyntaxError as e:
|
|
318
|
-
self.logger.warning(f"Failed to parse expression '{expr_str}': {e}")
|
|
319
|
-
return None
|
|
320
|
-
new_values.append(ast.FormattedValue(
|
|
321
|
-
value=expr_ast,
|
|
322
|
-
conversion=-1,
|
|
323
|
-
format_spec=None
|
|
324
|
-
))
|
|
325
|
-
idx = match.end()
|
|
326
|
-
else:
|
|
327
|
-
new_values.append(ast.Constant(value=formatted_sql[idx:]))
|
|
328
|
-
break
|
|
329
|
-
|
|
330
|
-
new_node = ast.JoinedStr(values=new_values)
|
|
331
|
-
if '\n' in formatted_sql:
|
|
332
|
-
if in_function:
|
|
333
|
-
formatted_fstring = ast.JoinedStr(values=[
|
|
334
|
-
ast.Constant(value='\n'),
|
|
335
|
-
*new_node.values,
|
|
336
|
-
ast.Constant(value=f'\n{" " * 4}')
|
|
337
|
-
])
|
|
338
|
-
else:
|
|
339
|
-
formatted_fstring = ast.JoinedStr(values=[
|
|
340
|
-
ast.Constant(value='\n'),
|
|
341
|
-
*new_node.values,
|
|
342
|
-
ast.Constant(value='\n')
|
|
343
|
-
])
|
|
344
|
-
else:
|
|
345
|
-
formatted_fstring = new_node
|
|
346
|
-
return formatted_fstring
|
|
301
|
+
self.changed = True
|
|
302
|
+
return self._build_replacement_text(formatted_sql, is_fstring, in_function)
|
|
347
303
|
|
|
348
|
-
return None
|
|
349
304
|
except SQLFormattingError as e:
|
|
350
305
|
self.logger.warning(f"SQL formatting skipped: {e}")
|
|
351
306
|
return None
|
|
@@ -353,30 +308,54 @@ class SQLStringFormatter(ast.NodeTransformer):
|
|
|
353
308
|
self.logger.warning(f"Unexpected error during SQL formatting: {e}")
|
|
354
309
|
return None
|
|
355
310
|
|
|
356
|
-
def
|
|
311
|
+
def _build_replacement_text(
|
|
312
|
+
self, formatted_sql: str, is_fstring: bool, in_function: bool
|
|
313
|
+
) -> str:
|
|
314
|
+
"""Constructs the replacement string literal text."""
|
|
315
|
+
prefix = "f" if is_fstring else ""
|
|
316
|
+
if in_function and "\n" in formatted_sql:
|
|
317
|
+
indent = " " * 4
|
|
318
|
+
lines = formatted_sql.split("\n")
|
|
319
|
+
indented_lines = [indent + line if line.strip() else line for line in lines]
|
|
320
|
+
indented_sql = "\n".join(indented_lines)
|
|
321
|
+
return f'\n{indent}{prefix}"""\n{indented_sql}\n{indent}"""\n'
|
|
322
|
+
elif "\n" in formatted_sql:
|
|
323
|
+
return f'{prefix}"""\n{formatted_sql}\n"""'
|
|
324
|
+
else:
|
|
325
|
+
return f'{prefix}"{formatted_sql}"'
|
|
326
|
+
|
|
327
|
+
def visit_Assign(self, node: ast.Assign) -> None:
|
|
357
328
|
"""Handles assignments."""
|
|
358
329
|
if isinstance(node.value, (ast.Constant, ast.JoinedStr)):
|
|
359
|
-
|
|
360
|
-
if
|
|
361
|
-
node.value
|
|
362
|
-
|
|
330
|
+
replacement_text = self.format_sql_node(node.value)
|
|
331
|
+
if replacement_text:
|
|
332
|
+
start, end = node_offsets(node.value, self.line_offsets)
|
|
333
|
+
self.replacements.append(SQLReplacement(start, end, replacement_text))
|
|
334
|
+
self.generic_visit(node)
|
|
363
335
|
|
|
364
|
-
def visit_Call(self, node: ast.Call) ->
|
|
336
|
+
def visit_Call(self, node: ast.Call) -> None:
|
|
365
337
|
"""Handles function calls."""
|
|
366
338
|
func_name = self.get_full_func_name(node.func)
|
|
367
339
|
if any(name in func_name for name in self.config.function_names):
|
|
368
|
-
for
|
|
340
|
+
for arg in node.args:
|
|
369
341
|
if isinstance(arg, (ast.Constant, ast.JoinedStr)):
|
|
370
|
-
|
|
371
|
-
|
|
372
|
-
|
|
373
|
-
|
|
342
|
+
replacement_text = self.format_sql_node(arg, in_function=True)
|
|
343
|
+
if replacement_text:
|
|
344
|
+
start, end = node_offsets(arg, self.line_offsets)
|
|
345
|
+
self.replacements.append(
|
|
346
|
+
SQLReplacement(start, end, replacement_text)
|
|
347
|
+
)
|
|
374
348
|
for keyword in node.keywords:
|
|
375
349
|
if isinstance(keyword.value, (ast.Constant, ast.JoinedStr)):
|
|
376
|
-
|
|
377
|
-
|
|
378
|
-
|
|
379
|
-
|
|
350
|
+
replacement_text = self.format_sql_node(
|
|
351
|
+
keyword.value, in_function=True
|
|
352
|
+
)
|
|
353
|
+
if replacement_text:
|
|
354
|
+
start, end = node_offsets(keyword.value, self.line_offsets)
|
|
355
|
+
self.replacements.append(
|
|
356
|
+
SQLReplacement(start, end, replacement_text)
|
|
357
|
+
)
|
|
358
|
+
self.generic_visit(node)
|
|
380
359
|
|
|
381
360
|
@staticmethod
|
|
382
361
|
def get_full_func_name(node: ast.AST) -> str:
|
|
@@ -477,19 +456,26 @@ def process_notebook(
|
|
|
477
456
|
|
|
478
457
|
continue # Move to the next cell after handling magic command
|
|
479
458
|
|
|
480
|
-
# Handle regular Python code
|
|
459
|
+
# Handle regular Python code with surgical string replacement
|
|
481
460
|
try:
|
|
482
461
|
tree = ast.parse(original_code)
|
|
483
|
-
|
|
484
|
-
|
|
485
|
-
|
|
486
|
-
|
|
487
|
-
|
|
488
|
-
|
|
489
|
-
|
|
490
|
-
|
|
491
|
-
|
|
492
|
-
|
|
462
|
+
line_offsets = build_line_offsets(original_code)
|
|
463
|
+
formatter = SQLStringFormatter(config, dialect, logger, line_offsets)
|
|
464
|
+
formatter.visit(tree)
|
|
465
|
+
if formatter.replacements:
|
|
466
|
+
result = original_code
|
|
467
|
+
for rep in sorted(
|
|
468
|
+
formatter.replacements,
|
|
469
|
+
key=lambda r: r.start_offset,
|
|
470
|
+
reverse=True,
|
|
471
|
+
):
|
|
472
|
+
result = (
|
|
473
|
+
result[: rep.start_offset]
|
|
474
|
+
+ rep.new_text
|
|
475
|
+
+ result[rep.end_offset :]
|
|
476
|
+
)
|
|
477
|
+
if result != original_code:
|
|
478
|
+
cell.source = result
|
|
493
479
|
changed = True
|
|
494
480
|
except SyntaxError:
|
|
495
481
|
logger.warning(f"Failed to parse cell:\n{original_code}")
|
|
@@ -556,4 +542,4 @@ def main():
|
|
|
556
542
|
|
|
557
543
|
|
|
558
544
|
if __name__ == "__main__":
|
|
559
|
-
main()
|
|
545
|
+
main()
|
|
@@ -1,19 +1,19 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: sqlnbfmt
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
Summary: A tool to format SQL code in Jupyter notebooks.
|
|
5
5
|
Author-email: Qing <qingye779@gmail.com>
|
|
6
6
|
License: MIT
|
|
7
7
|
Classifier: Programming Language :: Python :: 3
|
|
8
8
|
Classifier: Programming Language :: Python :: 3.12
|
|
9
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
9
10
|
Requires-Python: >=3.12
|
|
10
11
|
Description-Content-Type: text/markdown
|
|
11
12
|
License-File: LICENSE
|
|
12
|
-
Requires-Dist: astor>=0.8.1
|
|
13
|
-
Requires-Dist: black>=24.10.0
|
|
14
13
|
Requires-Dist: nbformat>=5.10.4
|
|
15
14
|
Requires-Dist: pyyaml>=6.0.2
|
|
16
15
|
Requires-Dist: sqlglot>=25.27.0
|
|
16
|
+
Dynamic: license-file
|
|
17
17
|
|
|
18
18
|
# sqlnbfmt
|
|
19
19
|
|
|
@@ -27,10 +27,10 @@ A SQL formatter designed specifically for Jupyter Notebooks. `sqlnbfmt` automati
|
|
|
27
27
|
|
|
28
28
|
- 🎯 **Smart SQL Detection**: Automatically identifies and formats SQL queries in code cells and magic SQL cells
|
|
29
29
|
- 🌳 **AST-Powered**: Uses Abstract Syntax Tree parsing for accurate SQL string identification
|
|
30
|
-
- 🔒 **Safe Formatting**: Preserves query parameters (e.g., `%s`, `?`) and SQL comments
|
|
30
|
+
- 🔒 **Safe Formatting**: Preserves Python comments, query parameters (e.g., `%s`, `?`), and SQL comments
|
|
31
31
|
- ⚙️ **Highly Configurable**: Customize formatting through YAML configuration
|
|
32
32
|
- 🔄 **Pre-commit Ready**: Seamlessly integrates with pre-commit hooks
|
|
33
|
-
- 📦 **
|
|
33
|
+
- 📦 **Lightweight**: Only three runtime dependencies (sqlglot, nbformat, pyyaml)
|
|
34
34
|
|
|
35
35
|
## Installation
|
|
36
36
|
|
|
@@ -63,14 +63,14 @@ pip install pre-commit
|
|
|
63
63
|
```yaml
|
|
64
64
|
repos:
|
|
65
65
|
- repo: https://github.com/flyersworder/sqlnbfmt
|
|
66
|
-
rev: v0.
|
|
66
|
+
rev: v0.2.0
|
|
67
67
|
hooks:
|
|
68
68
|
- id: sqlnbfmt
|
|
69
69
|
name: sqlnbfmt
|
|
70
70
|
types: [jupyter]
|
|
71
71
|
args: [--config, config.yaml, --dialect, postgres]
|
|
72
72
|
```
|
|
73
|
-
Please
|
|
73
|
+
Please run the following command for help:
|
|
74
74
|
|
|
75
75
|
```bash
|
|
76
76
|
sqlnbfmt -h
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
"""Fixture-based eval test suite for sqlnbfmt.
|
|
2
|
+
|
|
3
|
+
Auto-discovers test cases from tests/eval/*/ directories.
|
|
4
|
+
Each directory must contain input.ipynb and expected.ipynb.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import logging
|
|
8
|
+
import shutil
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
|
|
11
|
+
import nbformat
|
|
12
|
+
import pytest
|
|
13
|
+
|
|
14
|
+
from sqlnbfmt.formatter import load_config, process_notebook
|
|
15
|
+
|
|
16
|
+
EVAL_DIR = Path(__file__).resolve().parent / "eval"
|
|
17
|
+
CONFIG_PATH = Path(__file__).resolve().parent.parent / "config.yaml"
|
|
18
|
+
DIALECT = "mysql"
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def discover_eval_cases():
|
|
22
|
+
"""Find all eval case directories containing input.ipynb + expected.ipynb."""
|
|
23
|
+
if not EVAL_DIR.is_dir():
|
|
24
|
+
return []
|
|
25
|
+
cases = []
|
|
26
|
+
for case_dir in sorted(EVAL_DIR.iterdir()):
|
|
27
|
+
if (
|
|
28
|
+
case_dir.is_dir()
|
|
29
|
+
and (case_dir / "input.ipynb").exists()
|
|
30
|
+
and (case_dir / "expected.ipynb").exists()
|
|
31
|
+
):
|
|
32
|
+
cases.append(case_dir.name)
|
|
33
|
+
return cases
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
@pytest.fixture
|
|
37
|
+
def logger():
|
|
38
|
+
log = logging.getLogger("test_eval")
|
|
39
|
+
log.setLevel(logging.DEBUG)
|
|
40
|
+
if not log.handlers:
|
|
41
|
+
handler = logging.StreamHandler()
|
|
42
|
+
handler.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
|
|
43
|
+
log.addHandler(handler)
|
|
44
|
+
return log
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
@pytest.mark.parametrize("case_name", discover_eval_cases(), ids=lambda c: c)
|
|
48
|
+
def test_eval(case_name, logger, tmp_path):
|
|
49
|
+
"""Run formatter on input.ipynb and compare cell-by-cell with expected.ipynb."""
|
|
50
|
+
case_dir = EVAL_DIR / case_name
|
|
51
|
+
|
|
52
|
+
# Copy input to temp directory so we don't modify fixtures
|
|
53
|
+
work_path = tmp_path / "notebook.ipynb"
|
|
54
|
+
shutil.copy2(case_dir / "input.ipynb", work_path)
|
|
55
|
+
|
|
56
|
+
# Load config and run formatter
|
|
57
|
+
config = load_config(CONFIG_PATH)
|
|
58
|
+
process_notebook(work_path, config, DIALECT, logger)
|
|
59
|
+
|
|
60
|
+
# Read formatted and expected
|
|
61
|
+
formatted_nb = nbformat.read(work_path, as_version=4)
|
|
62
|
+
expected_nb = nbformat.read(case_dir / "expected.ipynb", as_version=4)
|
|
63
|
+
|
|
64
|
+
# Compare cell-by-cell
|
|
65
|
+
assert len(formatted_nb.cells) == len(expected_nb.cells), (
|
|
66
|
+
f"[{case_name}] Cell count mismatch: "
|
|
67
|
+
f"got {len(formatted_nb.cells)}, expected {len(expected_nb.cells)}"
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
for i, (fmt_cell, exp_cell) in enumerate(
|
|
71
|
+
zip(formatted_nb.cells, expected_nb.cells)
|
|
72
|
+
):
|
|
73
|
+
assert fmt_cell.cell_type == exp_cell.cell_type, (
|
|
74
|
+
f"[{case_name}] Cell {i} type mismatch: "
|
|
75
|
+
f"got {fmt_cell.cell_type}, expected {exp_cell.cell_type}"
|
|
76
|
+
)
|
|
77
|
+
assert fmt_cell.source == exp_cell.source, (
|
|
78
|
+
f"[{case_name}] Cell {i} content mismatch.\n"
|
|
79
|
+
f"Expected:\n{exp_cell.source}\n"
|
|
80
|
+
f"Actual:\n{fmt_cell.source}"
|
|
81
|
+
)
|
|
@@ -0,0 +1,237 @@
|
|
|
1
|
+
import pytest
|
|
2
|
+
import nbformat
|
|
3
|
+
from nbformat import v4 as nbf
|
|
4
|
+
import logging
|
|
5
|
+
|
|
6
|
+
from sqlnbfmt.formatter import process_notebook, load_config
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
# Fixture for temporary notebook path
|
|
10
|
+
@pytest.fixture
|
|
11
|
+
def temp_nb_path(tmp_path):
|
|
12
|
+
return tmp_path / "temp_test_notebook.ipynb"
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
# Fixture for logger
|
|
16
|
+
@pytest.fixture
|
|
17
|
+
def logger():
|
|
18
|
+
# Set up a logger for testing
|
|
19
|
+
logger = logging.getLogger("test_formatter")
|
|
20
|
+
logger.setLevel(logging.DEBUG)
|
|
21
|
+
if not logger.handlers:
|
|
22
|
+
handler = logging.StreamHandler()
|
|
23
|
+
formatter = logging.Formatter("[%(levelname)s] %(message)s")
|
|
24
|
+
handler.setFormatter(formatter)
|
|
25
|
+
logger.addHandler(handler)
|
|
26
|
+
return logger
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
# Test function with updated expected outputs
|
|
30
|
+
@pytest.mark.parametrize(
|
|
31
|
+
"input_cells, expected_cells",
|
|
32
|
+
[
|
|
33
|
+
# Test Case 1: Simple SELECT Query
|
|
34
|
+
(
|
|
35
|
+
[
|
|
36
|
+
nbf.new_code_cell(
|
|
37
|
+
"""query = "select id, name from users where active = 1" """
|
|
38
|
+
)
|
|
39
|
+
],
|
|
40
|
+
[
|
|
41
|
+
nbf.new_code_cell(
|
|
42
|
+
'''query = """\nSELECT\n id,\n name\nFROM users\nWHERE\n active = 1\n""" '''
|
|
43
|
+
)
|
|
44
|
+
],
|
|
45
|
+
),
|
|
46
|
+
# Test Case 2: SELECT Query with F-String and Variable
|
|
47
|
+
(
|
|
48
|
+
[
|
|
49
|
+
nbf.new_code_cell(
|
|
50
|
+
'''table_name = "users"\nquery = f"select * from {table_name}"'''
|
|
51
|
+
)
|
|
52
|
+
],
|
|
53
|
+
[
|
|
54
|
+
nbf.new_code_cell(
|
|
55
|
+
'''table_name = "users"\nquery = f"""\nSELECT\n *\nFROM {table_name}\n"""'''
|
|
56
|
+
)
|
|
57
|
+
],
|
|
58
|
+
),
|
|
59
|
+
# Test Case 3: Multi-line SELECT Query with Variables
|
|
60
|
+
(
|
|
61
|
+
[
|
|
62
|
+
nbf.new_code_cell(
|
|
63
|
+
'''columns = "id, name"\nquery = f"select {columns} from users where active = 1"'''
|
|
64
|
+
)
|
|
65
|
+
],
|
|
66
|
+
[
|
|
67
|
+
nbf.new_code_cell(
|
|
68
|
+
'''columns = "id, name"\nquery = f"""\nSELECT\n {columns}\nFROM users\nWHERE\n active = 1\n"""'''
|
|
69
|
+
)
|
|
70
|
+
],
|
|
71
|
+
),
|
|
72
|
+
# Test Case 4: SQL Query with Line Magic Command
|
|
73
|
+
(
|
|
74
|
+
[nbf.new_code_cell("""%sql select * from users where active = 1""")],
|
|
75
|
+
[nbf.new_code_cell("""%sql SELECT * FROM users WHERE active = 1""")],
|
|
76
|
+
),
|
|
77
|
+
# Test Case 5: SQL Query with Cell Magic Command
|
|
78
|
+
(
|
|
79
|
+
[nbf.new_code_cell("""%%sql\nselect * from users where active = 1""")],
|
|
80
|
+
[
|
|
81
|
+
nbf.new_code_cell(
|
|
82
|
+
"""%%sql\nSELECT\n *\nFROM users\nWHERE\n active = 1"""
|
|
83
|
+
)
|
|
84
|
+
],
|
|
85
|
+
),
|
|
86
|
+
# Test Case 6: INSERT Query with Variables (Remains Unchanged)
|
|
87
|
+
(
|
|
88
|
+
[
|
|
89
|
+
nbf.new_code_cell(
|
|
90
|
+
'''table_name = "users"\ncolumns = "(id, name)"\nvalues = "(1, 'Alice')"\nquery = f"insert into {table_name} {columns} values {values}"'''
|
|
91
|
+
)
|
|
92
|
+
],
|
|
93
|
+
[
|
|
94
|
+
nbf.new_code_cell(
|
|
95
|
+
'''table_name = "users"\ncolumns = "(id, name)"\nvalues = "(1, 'Alice')"\nquery = f"insert into {table_name} {columns} values {values}"'''
|
|
96
|
+
)
|
|
97
|
+
],
|
|
98
|
+
),
|
|
99
|
+
# Test Case 7: UPDATE Query with Variables
|
|
100
|
+
(
|
|
101
|
+
[
|
|
102
|
+
nbf.new_code_cell(
|
|
103
|
+
'''table_name = "users"\nquery = f"update {table_name} set active = 0 where id = 1"'''
|
|
104
|
+
)
|
|
105
|
+
],
|
|
106
|
+
[
|
|
107
|
+
nbf.new_code_cell(
|
|
108
|
+
'''table_name = "users"\nquery = f"""\nUPDATE {table_name} SET active = 0\nWHERE\n id = 1\n"""'''
|
|
109
|
+
)
|
|
110
|
+
],
|
|
111
|
+
),
|
|
112
|
+
# Test Case 8: DELETE Query
|
|
113
|
+
(
|
|
114
|
+
[nbf.new_code_cell("""query = "delete from users where active = 0" """)],
|
|
115
|
+
[
|
|
116
|
+
nbf.new_code_cell(
|
|
117
|
+
'''query = """\nDELETE FROM users\nWHERE\n active = 0\n""" '''
|
|
118
|
+
)
|
|
119
|
+
],
|
|
120
|
+
),
|
|
121
|
+
# Test Case 9: Query with SQL Comments
|
|
122
|
+
(
|
|
123
|
+
[
|
|
124
|
+
nbf.new_code_cell(
|
|
125
|
+
'''query = "-- Select active users\\nselect id, name from users where active = 1"'''
|
|
126
|
+
)
|
|
127
|
+
],
|
|
128
|
+
[
|
|
129
|
+
nbf.new_code_cell(
|
|
130
|
+
'''query = """\n/* Select active users */\nSELECT\n id,\n name\nFROM users\nWHERE\n active = 1\n"""'''
|
|
131
|
+
)
|
|
132
|
+
],
|
|
133
|
+
),
|
|
134
|
+
# Test Case 10: Non-SQL String (Should Remain Unchanged)
|
|
135
|
+
(
|
|
136
|
+
[
|
|
137
|
+
nbf.new_code_cell(
|
|
138
|
+
'''message = "This is not an SQL query: select * from users;"'''
|
|
139
|
+
)
|
|
140
|
+
],
|
|
141
|
+
[
|
|
142
|
+
nbf.new_code_cell(
|
|
143
|
+
'''message = "This is not an SQL query: select * from users;"'''
|
|
144
|
+
)
|
|
145
|
+
],
|
|
146
|
+
),
|
|
147
|
+
# Test Case 11: F-string without variables
|
|
148
|
+
(
|
|
149
|
+
[
|
|
150
|
+
nbf.new_code_cell(
|
|
151
|
+
'''query = f"select id, name from users where active = 1"'''
|
|
152
|
+
)
|
|
153
|
+
],
|
|
154
|
+
[
|
|
155
|
+
nbf.new_code_cell(
|
|
156
|
+
'''query = f"""\nSELECT\n id,\n name\nFROM users\nWHERE\n active = 1\n"""'''
|
|
157
|
+
)
|
|
158
|
+
],
|
|
159
|
+
),
|
|
160
|
+
# Test Case 12: F-string with only string constants
|
|
161
|
+
(
|
|
162
|
+
[nbf.new_code_cell('''query = f"select " f"id, name " f"from users"''')],
|
|
163
|
+
[
|
|
164
|
+
nbf.new_code_cell(
|
|
165
|
+
'''query = f"""\nSELECT\n id,\n name\nFROM users\n"""'''
|
|
166
|
+
)
|
|
167
|
+
],
|
|
168
|
+
),
|
|
169
|
+
# Test Case 13: In-function SQL Query
|
|
170
|
+
(
|
|
171
|
+
[
|
|
172
|
+
nbf.new_code_cell(
|
|
173
|
+
'''pd.read_sql("""select id, name from users where active = 1""")'''
|
|
174
|
+
)
|
|
175
|
+
],
|
|
176
|
+
[
|
|
177
|
+
nbf.new_code_cell(
|
|
178
|
+
'''pd.read_sql(\n """\n SELECT\n id,\n name\n FROM users\n WHERE\n active = 1\n """\n)'''
|
|
179
|
+
)
|
|
180
|
+
],
|
|
181
|
+
),
|
|
182
|
+
# Test Case 14: In-function SQL Query with F-String
|
|
183
|
+
(
|
|
184
|
+
[
|
|
185
|
+
nbf.new_code_cell(
|
|
186
|
+
"""table_name = "users"\npd.read_sql(f"select * from {table_name}")"""
|
|
187
|
+
)
|
|
188
|
+
],
|
|
189
|
+
[
|
|
190
|
+
nbf.new_code_cell(
|
|
191
|
+
'''table_name = "users"\npd.read_sql(\n f"""\n SELECT\n *\n FROM {table_name}\n """\n)'''
|
|
192
|
+
)
|
|
193
|
+
],
|
|
194
|
+
),
|
|
195
|
+
# Test Case 15: Python comment preservation
|
|
196
|
+
(
|
|
197
|
+
[
|
|
198
|
+
nbf.new_code_cell(
|
|
199
|
+
'''# This query fetches active users\nquery = "select id, name from users where active = 1"'''
|
|
200
|
+
)
|
|
201
|
+
],
|
|
202
|
+
[
|
|
203
|
+
nbf.new_code_cell(
|
|
204
|
+
'''# This query fetches active users\nquery = """\nSELECT\n id,\n name\nFROM users\nWHERE\n active = 1\n"""'''
|
|
205
|
+
)
|
|
206
|
+
],
|
|
207
|
+
),
|
|
208
|
+
],
|
|
209
|
+
)
|
|
210
|
+
def test_sql_formatter(temp_nb_path, logger, input_cells, expected_cells):
|
|
211
|
+
# Create a new notebook
|
|
212
|
+
nb = nbf.new_notebook()
|
|
213
|
+
nb.cells = input_cells
|
|
214
|
+
|
|
215
|
+
# Write the notebook to the temporary path
|
|
216
|
+
nbformat.write(nb, temp_nb_path)
|
|
217
|
+
|
|
218
|
+
# Load configuration
|
|
219
|
+
config = load_config("config.yaml")
|
|
220
|
+
dialect = "mysql" # or your preferred SQL dialect
|
|
221
|
+
|
|
222
|
+
# Run the formatter on the notebook
|
|
223
|
+
process_notebook(temp_nb_path, config, dialect, logger=logger)
|
|
224
|
+
|
|
225
|
+
# Read the formatted notebook
|
|
226
|
+
formatted_nb = nbformat.read(temp_nb_path, as_version=4)
|
|
227
|
+
|
|
228
|
+
# Compare the cells
|
|
229
|
+
for formatted_cell, expected_cell in zip(formatted_nb.cells, expected_cells):
|
|
230
|
+
formatted_code = formatted_cell.source.strip()
|
|
231
|
+
expected_code = expected_cell.source.strip()
|
|
232
|
+
assert formatted_code == expected_code, (
|
|
233
|
+
f"Formatted code does not match expected output.\nExpected:\n{expected_code}\nActual:\n{formatted_code}"
|
|
234
|
+
)
|
|
235
|
+
|
|
236
|
+
# Clean up
|
|
237
|
+
temp_nb_path.unlink()
|
|
@@ -1,201 +0,0 @@
|
|
|
1
|
-
import pytest
|
|
2
|
-
import nbformat
|
|
3
|
-
from nbformat import v4 as nbf
|
|
4
|
-
from pathlib import Path
|
|
5
|
-
import logging
|
|
6
|
-
import sys
|
|
7
|
-
|
|
8
|
-
# Add the parent directory to sys.path to import the formatter module
|
|
9
|
-
sys.path.append(str(Path(__file__).resolve().parent.parent))
|
|
10
|
-
|
|
11
|
-
from sqlnbfmt.formatter import process_notebook, FormattingConfig, load_config
|
|
12
|
-
|
|
13
|
-
# Fixture for temporary notebook path
|
|
14
|
-
@pytest.fixture
|
|
15
|
-
def temp_nb_path(tmp_path):
|
|
16
|
-
return tmp_path / "temp_test_notebook.ipynb"
|
|
17
|
-
|
|
18
|
-
# Fixture for logger
|
|
19
|
-
@pytest.fixture
|
|
20
|
-
def logger():
|
|
21
|
-
# Set up a logger for testing
|
|
22
|
-
logger = logging.getLogger('test_formatter')
|
|
23
|
-
logger.setLevel(logging.DEBUG)
|
|
24
|
-
if not logger.handlers:
|
|
25
|
-
handler = logging.StreamHandler()
|
|
26
|
-
formatter = logging.Formatter("[%(levelname)s] %(message)s")
|
|
27
|
-
handler.setFormatter(formatter)
|
|
28
|
-
logger.addHandler(handler)
|
|
29
|
-
return logger
|
|
30
|
-
|
|
31
|
-
# Test function with updated expected outputs
|
|
32
|
-
@pytest.mark.parametrize("input_cells, expected_cells", [
|
|
33
|
-
# Test Case 1: Simple SELECT Query
|
|
34
|
-
(
|
|
35
|
-
[nbf.new_code_cell("""query = "select id, name from users where active = 1" """)],
|
|
36
|
-
[nbf.new_code_cell('''query = """
|
|
37
|
-
SELECT
|
|
38
|
-
id,
|
|
39
|
-
name
|
|
40
|
-
FROM users
|
|
41
|
-
WHERE
|
|
42
|
-
active = 1
|
|
43
|
-
"""''')],
|
|
44
|
-
),
|
|
45
|
-
# Test Case 2: SELECT Query with F-String and Variable
|
|
46
|
-
(
|
|
47
|
-
[nbf.new_code_cell('''table_name = "users"\nquery = f"select * from {table_name}"''')],
|
|
48
|
-
[nbf.new_code_cell('''table_name = "users"
|
|
49
|
-
query = f"""
|
|
50
|
-
SELECT
|
|
51
|
-
*
|
|
52
|
-
FROM {table_name}
|
|
53
|
-
"""''')],
|
|
54
|
-
),
|
|
55
|
-
# Test Case 3: Multi-line SELECT Query with Variables
|
|
56
|
-
(
|
|
57
|
-
[nbf.new_code_cell('''columns = "id, name"\nquery = f"select {columns} from users where active = 1"''')],
|
|
58
|
-
[nbf.new_code_cell('''columns = "id, name"
|
|
59
|
-
query = f"""
|
|
60
|
-
SELECT
|
|
61
|
-
{columns}
|
|
62
|
-
FROM users
|
|
63
|
-
WHERE
|
|
64
|
-
active = 1
|
|
65
|
-
"""''')],
|
|
66
|
-
),
|
|
67
|
-
# Test Case 4: SQL Query with Line Magic Command
|
|
68
|
-
(
|
|
69
|
-
[nbf.new_code_cell("""%sql select * from users where active = 1""")],
|
|
70
|
-
[nbf.new_code_cell("""%sql SELECT * FROM users WHERE active = 1""")],
|
|
71
|
-
),
|
|
72
|
-
# Test Case 5: SQL Query with Cell Magic Command
|
|
73
|
-
(
|
|
74
|
-
[nbf.new_code_cell('''%%sql\nselect * from users where active = 1''')],
|
|
75
|
-
[nbf.new_code_cell('''%%sql
|
|
76
|
-
SELECT
|
|
77
|
-
*
|
|
78
|
-
FROM users
|
|
79
|
-
WHERE
|
|
80
|
-
active = 1''')],
|
|
81
|
-
),
|
|
82
|
-
# Test Case 6: INSERT Query with Variables (Remains Unchanged)
|
|
83
|
-
(
|
|
84
|
-
[nbf.new_code_cell('''table_name = "users"\ncolumns = "(id, name)"\nvalues = "(1, 'Alice')"\nquery = f"insert into {table_name} {columns} values {values}"''')],
|
|
85
|
-
[nbf.new_code_cell('''table_name = "users"
|
|
86
|
-
columns = "(id, name)"
|
|
87
|
-
values = "(1, 'Alice')"
|
|
88
|
-
query = f"insert into {table_name} {columns} values {values}"''')],
|
|
89
|
-
),
|
|
90
|
-
# Test Case 7: UPDATE Query with Variables
|
|
91
|
-
(
|
|
92
|
-
[nbf.new_code_cell('''table_name = "users"\nquery = f"update {table_name} set active = 0 where id = 1"''')],
|
|
93
|
-
[nbf.new_code_cell('''table_name = "users"
|
|
94
|
-
query = f"""
|
|
95
|
-
UPDATE {table_name} SET active = 0
|
|
96
|
-
WHERE
|
|
97
|
-
id = 1
|
|
98
|
-
"""''')],
|
|
99
|
-
),
|
|
100
|
-
# Test Case 8: DELETE Query
|
|
101
|
-
(
|
|
102
|
-
[nbf.new_code_cell("""query = "delete from users where active = 0" """)],
|
|
103
|
-
[nbf.new_code_cell('''query = """
|
|
104
|
-
DELETE FROM users
|
|
105
|
-
WHERE
|
|
106
|
-
active = 0
|
|
107
|
-
"""''')],
|
|
108
|
-
),
|
|
109
|
-
# Test Case 9: Query with Comments
|
|
110
|
-
(
|
|
111
|
-
[nbf.new_code_cell('''query = "-- Select active users\\nselect id, name from users where active = 1"''')],
|
|
112
|
-
[nbf.new_code_cell('''query = """
|
|
113
|
-
/* Select active users */
|
|
114
|
-
SELECT
|
|
115
|
-
id,
|
|
116
|
-
name
|
|
117
|
-
FROM users
|
|
118
|
-
WHERE
|
|
119
|
-
active = 1
|
|
120
|
-
"""''')],
|
|
121
|
-
),
|
|
122
|
-
# Test Case 10: Non-SQL String (Should Remain Unchanged)
|
|
123
|
-
(
|
|
124
|
-
[nbf.new_code_cell('''message = "This is not an SQL query: select * from users;"''')],
|
|
125
|
-
[nbf.new_code_cell('''message = "This is not an SQL query: select * from users;"''')],
|
|
126
|
-
),
|
|
127
|
-
# Test Case 11: F-string without variables
|
|
128
|
-
(
|
|
129
|
-
[nbf.new_code_cell('''query = f"select id, name from users where active = 1"''')],
|
|
130
|
-
[nbf.new_code_cell('''query = f"""
|
|
131
|
-
SELECT
|
|
132
|
-
id,
|
|
133
|
-
name
|
|
134
|
-
FROM users
|
|
135
|
-
WHERE
|
|
136
|
-
active = 1
|
|
137
|
-
"""''')],
|
|
138
|
-
),
|
|
139
|
-
# Test Case 12: F-string with only string constants
|
|
140
|
-
(
|
|
141
|
-
[nbf.new_code_cell('''query = f"select " f"id, name " f"from users"''')],
|
|
142
|
-
[nbf.new_code_cell('''query = f"""
|
|
143
|
-
SELECT
|
|
144
|
-
id,
|
|
145
|
-
name
|
|
146
|
-
FROM users
|
|
147
|
-
"""''')],
|
|
148
|
-
),
|
|
149
|
-
# Test Case 13: In-function SQL Query
|
|
150
|
-
(
|
|
151
|
-
[nbf.new_code_cell('''pd.read_sql("""select id, name from users where active = 1""")''')],
|
|
152
|
-
[nbf.new_code_cell('''pd.read_sql(
|
|
153
|
-
"""
|
|
154
|
-
SELECT
|
|
155
|
-
id,
|
|
156
|
-
name
|
|
157
|
-
FROM users
|
|
158
|
-
WHERE
|
|
159
|
-
active = 1
|
|
160
|
-
"""
|
|
161
|
-
)''')],
|
|
162
|
-
),
|
|
163
|
-
# Test Case 14: In-function SQL Query with F-String
|
|
164
|
-
(
|
|
165
|
-
[nbf.new_code_cell('''table_name = "users"\npd.read_sql(f"select * from {table_name}")''')],
|
|
166
|
-
[nbf.new_code_cell('''table_name = "users"
|
|
167
|
-
pd.read_sql(
|
|
168
|
-
f"""
|
|
169
|
-
SELECT
|
|
170
|
-
*
|
|
171
|
-
FROM {table_name}
|
|
172
|
-
"""
|
|
173
|
-
)''')],
|
|
174
|
-
),
|
|
175
|
-
])
|
|
176
|
-
def test_sql_formatter(temp_nb_path, logger, input_cells, expected_cells):
|
|
177
|
-
# Create a new notebook
|
|
178
|
-
nb = nbf.new_notebook()
|
|
179
|
-
nb.cells = input_cells
|
|
180
|
-
|
|
181
|
-
# Write the notebook to the temporary path
|
|
182
|
-
nbformat.write(nb, temp_nb_path)
|
|
183
|
-
|
|
184
|
-
# Load configuration
|
|
185
|
-
config = load_config('config.yaml')
|
|
186
|
-
dialect = 'mysql' # or your preferred SQL dialect
|
|
187
|
-
|
|
188
|
-
# Run the formatter on the notebook
|
|
189
|
-
process_notebook(temp_nb_path, config, dialect, logger=logger)
|
|
190
|
-
|
|
191
|
-
# Read the formatted notebook
|
|
192
|
-
formatted_nb = nbformat.read(temp_nb_path, as_version=4)
|
|
193
|
-
|
|
194
|
-
# Compare the cells
|
|
195
|
-
for formatted_cell, expected_cell in zip(formatted_nb.cells, expected_cells):
|
|
196
|
-
formatted_code = formatted_cell.source.strip()
|
|
197
|
-
expected_code = expected_cell.source.strip()
|
|
198
|
-
assert formatted_code == expected_code, f"Formatted code does not match expected output.\nExpected:\n{expected_code}\nActual:\n{formatted_code}"
|
|
199
|
-
|
|
200
|
-
# Clean up
|
|
201
|
-
temp_nb_path.unlink()
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|