sqlnbfmt 0.1.9__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.
@@ -1,19 +1,19 @@
1
- Metadata-Version: 2.2
1
+ Metadata-Version: 2.4
2
2
  Name: sqlnbfmt
3
- Version: 0.1.9
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
- - 📦 **Zero Dependencies**: Minimal installation footprint
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.1.1
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 fun the following command in your CMD for help:
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
- - 📦 **Zero Dependencies**: Minimal installation footprint
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.1.1
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 fun the following command in your CMD for help:
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.1.9"
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
- [tool.uv]
24
- dev-dependencies = [
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('formatter')
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('formatter')
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'%s|\?')
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(f"'{temp_placeholder}'", temp_placeholder)
147
- formatted_sql = formatted_sql.replace(f'"{temp_placeholder}"', temp_placeholder)
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(temp_placeholder, original_placeholder)
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(temp_placeholder, original_placeholder)
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.NodeTransformer):
168
- """AST NodeTransformer that formats SQL strings."""
192
+ class SQLStringFormatter(ast.NodeVisitor):
193
+ """AST NodeVisitor that collects SQL string replacements."""
169
194
 
170
195
  def __init__(
171
- self, config: FormattingConfig, dialect: Optional[str], logger: logging.Logger
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 = astor.to_source(value.value).strip()
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[ast.AST]:
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
- node.value,
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
- if formatted_sql != node.value:
246
- self.changed = True
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
- elif isinstance(node, ast.JoinedStr):
267
- sql_str, placeholders = self.extract_fstring_parts(node)
268
- if not sql_str or not self.is_likely_sql(sql_str):
269
- return None
270
-
271
- try:
272
- formatted_sql = format_sql_code(
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
- if placeholders:
283
- for placeholder in placeholders.keys():
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 visit_Assign(self, node: ast.Assign) -> ast.AST:
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
- formatted_node = self.format_sql_node(node.value)
360
- if formatted_node:
361
- node.value = formatted_node
362
- return self.generic_visit(node)
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) -> ast.AST:
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 idx, arg in enumerate(node.args):
340
+ for arg in node.args:
369
341
  if isinstance(arg, (ast.Constant, ast.JoinedStr)):
370
- # Format with Black-style indentation for function arguments
371
- formatted_node = self.format_sql_node(arg, in_function=True)
372
- if formatted_node:
373
- node.args[idx] = formatted_node
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
- formatted_node = self.format_sql_node(keyword.value, in_function=True)
377
- if formatted_node:
378
- keyword.value = formatted_node
379
- return self.generic_visit(node)
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
- formatter = SQLStringFormatter(config, dialect, logger)
484
- new_tree = formatter.visit(tree)
485
- if formatter.changed:
486
- # Use astor to convert AST back to source code
487
- formatted_code = astor.to_source(new_tree)
488
- # Now format with black
489
- formatted_code = black.format_str(formatted_code, mode=black.FileMode())
490
-
491
- if formatted_code != original_code:
492
- cell.source = formatted_code
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.2
1
+ Metadata-Version: 2.4
2
2
  Name: sqlnbfmt
3
- Version: 0.1.9
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
- - 📦 **Zero Dependencies**: Minimal installation footprint
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.1.1
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 fun the following command in your CMD for help:
73
+ Please run the following command for help:
74
74
 
75
75
  ```bash
76
76
  sqlnbfmt -h
@@ -9,4 +9,5 @@ sqlnbfmt.egg-info/dependency_links.txt
9
9
  sqlnbfmt.egg-info/entry_points.txt
10
10
  sqlnbfmt.egg-info/requires.txt
11
11
  sqlnbfmt.egg-info/top_level.txt
12
+ tests/test_eval.py
12
13
  tests/test_formatter.py
@@ -1,5 +1,3 @@
1
- astor>=0.8.1
2
- black>=24.10.0
3
1
  nbformat>=5.10.4
4
2
  pyyaml>=6.0.2
5
3
  sqlglot>=25.27.0
@@ -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