sqlnbfmt 0.3.0__tar.gz → 0.3.2__tar.gz

{sqlnbfmt-0.3.0/sqlnbfmt.egg-info → sqlnbfmt-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlnbfmt
-Version: 0.3.0
+Version: 0.3.2
 Summary: A tool to format SQL code in Jupyter notebooks.
 Author-email: Qing <qingye779@gmail.com>
 License: MIT

{sqlnbfmt-0.3.0 → sqlnbfmt-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlnbfmt"
-version = "0.3.0"
+version = "0.3.2"
 description = "A tool to format SQL code in Jupyter notebooks."
 authors = [
   { name = "Qing", email = "qingye779@gmail.com" }

{sqlnbfmt-0.3.0 → sqlnbfmt-0.3.2}/sqlnbfmt/formatter.py

@@ -10,7 +10,7 @@ from typing import Set, Optional, Dict, Union, Tuple
 
 import nbformat
 import yaml
-from sqlglot import parse_one, errors
+from sqlglot import parse, errors
 
 DEFAULT_SQL_KEYWORDS = frozenset(
     {
@@ -116,10 +116,16 @@ def load_config(config_path: Optional[Union[str, Path]] = None) -> FormattingCon
     try:
         with open(config_file) as file:
             config = yaml.safe_load(file)
+            if config is None:
+                return FormattingConfig()
             return FormattingConfig(
-                sql_keywords=set(config.get("sql_keywords", [])),
-                function_names=set(config.get("function_names", [])),
-                sql_decorators=set(config.get("sql_decorators", [])),
+                sql_keywords=set(config.get("sql_keywords", DEFAULT_SQL_KEYWORDS)),
+                function_names=set(
+                    config.get("function_names", DEFAULT_FUNCTION_NAMES)
+                ),
+                sql_decorators=set(
+                    config.get("sql_decorators", DEFAULT_SQL_DECORATORS)
+                ),
                 single_line_threshold=config.get("formatting_options", {}).get(
                     "single_line_threshold", 80
                 ),
@@ -200,10 +206,18 @@ def format_sql_code(
         temp_sql = temp_sql.strip()
 
         # Parse and format SQL
-        parsed = parse_one(temp_sql, read=dialect)
-        formatted_sql = parsed.sql(
-            pretty=not force_single_line, indent=config.indent_width, dialect=dialect
-        )
+        statements = parse(temp_sql, read=dialect)
+        formatted_parts = []
+        for stmt in statements:
+            formatted_parts.append(
+                stmt.sql(
+                    pretty=not force_single_line,
+                    indent=config.indent_width,
+                    dialect=dialect,
+                )
+            )
+        separator = "; " if force_single_line else ";\n"
+        formatted_sql = separator.join(formatted_parts)
 
         # Apply formatting based on context
         if is_magic_command and not is_cell_magic:
@@ -372,7 +386,7 @@ class SQLStringFormatter(ast.NodeVisitor):
         """Constructs the replacement string literal text."""
         prefix = "f" if is_fstring else ""
         if in_function and "\n" in formatted_sql:
-            indent = " " * 4
+            indent = " " * self.config.indent_width
             lines = formatted_sql.split("\n")
             indented_lines = [indent + line if line.strip() else line for line in lines]
             indented_sql = "\n".join(indented_lines)
@@ -427,188 +441,75 @@ class SQLStringFormatter(ast.NodeVisitor):
         return ".".join(reversed(parts))
 
 
-def process_notebook(
-    notebook_path: Union[str, Path],
-    config: FormattingConfig,
-    dialect: Optional[str],
-    logger: logging.Logger,
-    check_only: bool = False,
-) -> bool:
-    """Processes a Jupyter notebook.
+def _has_skip_hint(source: str) -> bool:
+    """Check if any comment line contains the sqlnbfmt: skip directive.
 
-    Args:
-        check_only: If True, detect changes but don't write the file.
-
-    Returns:
-        True if changes were made (or would be made in check mode).
+    Only dedicated comment lines are recognised (the line must start with ``#``
+    after stripping whitespace).  Inline trailing comments like
+    ``x = 1  # sqlnbfmt: skip`` are intentionally ignored so that the directive
+    cannot be triggered accidentally from within an expression or string.
     """
-    try:
-        notebook = nbformat.read(notebook_path, as_version=4)
-        changed = False
-
-        for cell_index, cell in enumerate(notebook.cells):
-            if cell.cell_type != "code":
-                continue
-
-            original_code = cell.source
-            if not original_code.strip():
-                continue
-
-            if "# sqlnbfmt: skip" in original_code:
-                logger.debug(f"[Cell {cell_index}] Skipping cell (sqlnbfmt: skip)")
-                continue
-
-            lines = original_code.split("\n")
-
-            # Initialize variables to track magic commands
-            magic_cmd = None
-            magic_cmd_index = None
-
-            # Iterate through lines to find the first non-comment magic command
-            for idx, line in enumerate(lines):
-                stripped = line.strip()
-
-                if not stripped:
-                    continue  # Skip empty lines
-
-                if stripped.startswith("#"):
-                    continue  # Skip comment lines
-
-                if stripped.startswith("%%sql") or stripped.startswith("%sql"):
-                    magic_cmd = stripped.split()[0]
-                    magic_cmd_index = idx
-                    break  # Magic command found
-
-                else:
-                    break  # Non-magic, non-comment line found
-
-            if magic_cmd:
-                is_cell_magic = magic_cmd.startswith("%%sql")
-
-                if is_cell_magic:
-                    # Cell magic: SQL code starts from the next line
-                    sql_code = "\n".join(lines[magic_cmd_index + 1 :]).strip()
-                else:
-                    # Line magic: SQL code is on the same line after the magic command
-                    sql_code = lines[magic_cmd_index][len(magic_cmd) :].strip()
-
-                try:
-                    formatted_sql = format_sql_code(
-                        sql_code,
-                        dialect,
-                        config,
-                        is_magic_command=True,
-                        is_cell_magic=is_cell_magic,
-                    )
-
-                    # Reconstruct the cell content
-                    if is_cell_magic:
-                        # Preserve comments before the magic command
-                        pre_magic = "\n".join(lines[:magic_cmd_index])
-                        if pre_magic:
-                            new_content = f"{pre_magic}\n{magic_cmd}\n{formatted_sql}"
-                        else:
-                            new_content = f"{magic_cmd}\n{formatted_sql}"
-                    else:
-                        # Preserve comments before the magic command
-                        pre_magic = "\n".join(lines[:magic_cmd_index])
-                        if pre_magic:
-                            new_content = f"{pre_magic}\n{magic_cmd} {formatted_sql}"
-                        else:
-                            new_content = f"{magic_cmd} {formatted_sql}"
-
-                    if new_content != original_code:
-                        cell.source = new_content
-                        changed = True
-
-                except SQLFormattingError as e:
-                    logger.warning(
-                        f"[Cell {cell_index}] SQL magic formatting skipped: {e}"
-                    )
-
-                continue  # Move to the next cell after handling magic command
-
-            # Handle regular Python code with surgical string replacement
-            try:
-                tree = ast.parse(original_code)
-                line_offsets = build_line_offsets(original_code)
-                formatter = SQLStringFormatter(config, dialect, logger, line_offsets)
-                formatter.visit(tree)
-                if formatter.replacements:
-                    result = original_code
-                    for rep in sorted(
-                        formatter.replacements,
-                        key=lambda r: r.start_offset,
-                        reverse=True,
-                    ):
-                        result = (
-                            result[: rep.start_offset]
-                            + rep.new_text
-                            + result[rep.end_offset :]
-                        )
-                    if result != original_code:
-                        cell.source = result
-                        changed = True
-            except SyntaxError:
-                logger.warning(
-                    f"[Cell {cell_index}] Failed to parse cell:\n{original_code}"
-                )
-                continue
-
-        if changed and not check_only:
-            nbformat.write(notebook, notebook_path)
-            logger.info(f"Updated notebook: {notebook_path}")
-
-        return changed
-
-    except Exception as e:
-        logger.error(f"Failed to process notebook {notebook_path}: {e}", exc_info=True)
-        raise
+    return any(
+        line.strip().startswith("#") and "sqlnbfmt: skip" in line
+        for line in source.split("\n")
+    )
 
 
-def diff_notebook(
-    notebook_path: Union[str, Path],
+def _format_cells(
+    notebook: nbformat.NotebookNode,
     config: FormattingConfig,
     dialect: Optional[str],
     logger: logging.Logger,
-) -> str:
-    """Returns a unified diff of formatting changes for a notebook."""
-    notebook_path = Path(notebook_path)
-    original_nb = nbformat.read(notebook_path, as_version=4)
-    formatted_nb = copy.deepcopy(original_nb)
+) -> bool:
+    """Format cells in-place. Returns True if any cell was modified."""
+    changed = False
 
-    # Format the copy in-memory by running the same cell logic
-    for cell_index, cell in enumerate(formatted_nb.cells):
+    for cell_index, cell in enumerate(notebook.cells):
         if cell.cell_type != "code":
             continue
+
         original_code = cell.source
         if not original_code.strip():
             continue
-        if "# sqlnbfmt: skip" in original_code:
+
+        if _has_skip_hint(original_code):
+            logger.debug(f"[Cell {cell_index}] Skipping cell (sqlnbfmt: skip)")
             continue
 
         lines = original_code.split("\n")
+
+        # Initialize variables to track magic commands
         magic_cmd = None
         magic_cmd_index = None
+
+        # Iterate through lines to find the first non-comment magic command
         for idx, line in enumerate(lines):
             stripped = line.strip()
+
             if not stripped:
-                continue
+                continue  # Skip empty lines
+
             if stripped.startswith("#"):
-                continue
+                continue  # Skip comment lines
+
             if stripped.startswith("%%sql") or stripped.startswith("%sql"):
                 magic_cmd = stripped.split()[0]
                 magic_cmd_index = idx
-                break
+                break  # Magic command found
+
             else:
-                break
+                break  # Non-magic, non-comment line found
 
         if magic_cmd:
             is_cell_magic = magic_cmd.startswith("%%sql")
+
             if is_cell_magic:
+                # Cell magic: SQL code starts from the next line
                 sql_code = "\n".join(lines[magic_cmd_index + 1 :]).strip()
             else:
+                # Line magic: SQL code is on the same line after the magic command
                 sql_code = lines[magic_cmd_index][len(magic_cmd) :].strip()
+
             try:
                 formatted_sql = format_sql_code(
                     sql_code,
@@ -617,23 +518,33 @@ def diff_notebook(
                     is_magic_command=True,
                     is_cell_magic=is_cell_magic,
                 )
+
+                # Reconstruct the cell content
                 if is_cell_magic:
+                    # Preserve comments before the magic command
                     pre_magic = "\n".join(lines[:magic_cmd_index])
                     if pre_magic:
                         new_content = f"{pre_magic}\n{magic_cmd}\n{formatted_sql}"
                     else:
                         new_content = f"{magic_cmd}\n{formatted_sql}"
                 else:
+                    # Preserve comments before the magic command
                     pre_magic = "\n".join(lines[:magic_cmd_index])
                     if pre_magic:
                         new_content = f"{pre_magic}\n{magic_cmd} {formatted_sql}"
                     else:
                         new_content = f"{magic_cmd} {formatted_sql}"
-                cell.source = new_content
-            except SQLFormattingError:
-                pass
-            continue
 
+                if new_content != original_code:
+                    cell.source = new_content
+                    changed = True
+
+            except SQLFormattingError as e:
+                logger.warning(f"[Cell {cell_index}] SQL magic formatting skipped: {e}")
+
+            continue  # Move to the next cell after handling magic command
+
+        # Handle regular Python code with surgical string replacement
         try:
             tree = ast.parse(original_code)
             line_offsets = build_line_offsets(original_code)
@@ -651,10 +562,61 @@ def diff_notebook(
                         + rep.new_text
                         + result[rep.end_offset :]
                     )
-                cell.source = result
+                if result != original_code:
+                    cell.source = result
+                    changed = True
         except SyntaxError:
+            logger.warning(
+                f"[Cell {cell_index}] Failed to parse cell:\n{original_code}"
+            )
             continue
 
+    return changed
+
+
+def process_notebook(
+    notebook_path: Union[str, Path],
+    config: FormattingConfig,
+    dialect: Optional[str],
+    logger: logging.Logger,
+    check_only: bool = False,
+) -> bool:
+    """Processes a Jupyter notebook.
+
+    Args:
+        check_only: If True, detect changes but don't write the file.
+
+    Returns:
+        True if changes were made (or would be made in check mode).
+    """
+    try:
+        notebook = nbformat.read(notebook_path, as_version=4)
+        changed = _format_cells(notebook, config, dialect, logger)
+
+        if changed and not check_only:
+            nbformat.write(notebook, notebook_path)
+            logger.info(f"Updated notebook: {notebook_path}")
+
+        return changed
+
+    except Exception as e:
+        logger.error(f"Failed to process notebook {notebook_path}: {e}", exc_info=True)
+        raise
+
+
+def diff_notebook(
+    notebook_path: Union[str, Path],
+    config: FormattingConfig,
+    dialect: Optional[str],
+    logger: logging.Logger,
+) -> str:
+    """Returns a unified diff of formatting changes for a notebook."""
+    notebook_path = Path(notebook_path)
+    original_nb = nbformat.read(notebook_path, as_version=4)
+    formatted_nb = copy.deepcopy(original_nb)
+
+    _format_cells(formatted_nb, config, dialect, logger)
+
     # Build unified diff cell-by-cell
     diff_lines = []
     for i, (orig_cell, fmt_cell) in enumerate(
@@ -665,8 +627,8 @@ def diff_notebook(
                 difflib.unified_diff(
                     orig_cell.source.splitlines(keepends=True),
                     fmt_cell.source.splitlines(keepends=True),
-                    fromfile=f"Cell [{i}] original",
-                    tofile=f"Cell [{i}] formatted",
+                    fromfile=f"{notebook_path} Cell [{i}] original",
+                    tofile=f"{notebook_path} Cell [{i}] formatted",
                 )
             )
     return "".join(diff_lines)
@@ -715,30 +677,27 @@ def main():
         sys.exit(1)
 
     try:
-        changed = False
         changed_files = []
 
         for notebook in args.notebooks:
-            if args.diff:
+            if args.check or args.diff:
+                # In check/diff mode, use diff_notebook to get both the
+                # diff output and whether changes are needed (single pass).
                 diff_output = diff_notebook(notebook, config, args.dialect, logger)
-                if diff_output:
+                needs_formatting = bool(diff_output)
+
+                if args.diff and diff_output:
                     print(diff_output, end="")
-                    changed = True
-                    changed_files.append(notebook)
 
-            if args.check:
-                if process_notebook(
-                    notebook, config, args.dialect, logger, check_only=True
-                ):
-                    changed = True
+                if needs_formatting:
                     changed_files.append(notebook)
-            elif not args.diff:
+            else:
+                # Normal mode: format in place
                 if process_notebook(notebook, config, args.dialect, logger):
-                    changed = True
                     changed_files.append(notebook)
 
         if args.check:
-            if changed:
+            if changed_files:
                 logger.error(f"{len(changed_files)} notebook(s) would be reformatted:")
                 for f in changed_files:
                     logger.error(f"  - {f}")
@@ -749,7 +708,7 @@ def main():
 
         # Print summary for normal mode
         if not args.diff:
-            if changed:
+            if changed_files:
                 logger.info("Changes made to the following notebooks:")
                 for file in changed_files:
                     logger.info(f"  - {file}")

{sqlnbfmt-0.3.0 → sqlnbfmt-0.3.2/sqlnbfmt.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlnbfmt
-Version: 0.3.0
+Version: 0.3.2
 Summary: A tool to format SQL code in Jupyter notebooks.
 Author-email: Qing <qingye779@gmail.com>
 License: MIT

{sqlnbfmt-0.3.0 → sqlnbfmt-0.3.2}/tests/test_formatter.py

@@ -334,3 +334,130 @@ def test_load_config_missing_file():
 
     with pytest.raises(FileNotFoundError):
         load_config("/nonexistent/config.yaml")
+
+
+def test_load_config_empty_file(tmp_path):
+    """load_config with empty YAML file returns defaults (not AttributeError)."""
+    empty_config = tmp_path / "empty.yaml"
+    empty_config.write_text("")
+    config = load_config(empty_config)
+    assert "SELECT" in config.sql_keywords
+    assert "read_sql" in config.function_names
+
+
+def test_load_config_partial_override(tmp_path):
+    """Config file with only indent_width still gets default keywords."""
+    partial_config = tmp_path / "partial.yaml"
+    partial_config.write_text("formatting_options:\n  indent_width: 2\n")
+    config = load_config(partial_config)
+    assert config.indent_width == 2
+    assert "SELECT" in config.sql_keywords
+    assert "read_sql" in config.function_names
+
+
+def test_skip_hint_in_string_not_triggered(tmp_path, logger):
+    """Skip hint inside a string literal should NOT skip the cell."""
+    nb = nbf.new_notebook()
+    original_source = (
+        'msg = "Use # sqlnbfmt: skip to disable"\n'
+        'query = "select id, name from users where active = 1"'
+    )
+    nb.cells = [nbf.new_code_cell(original_source)]
+    nb_path = tmp_path / "skip_in_string.ipynb"
+    nbformat.write(nb, nb_path)
+
+    config = load_config()
+    changed = process_notebook(nb_path, config, "mysql", logger)
+    assert changed  # SQL should still be formatted
+
+
+def test_skip_hint_magic_cell(tmp_path, logger):
+    """Skip hint in a magic command cell should skip formatting."""
+    nb = nbf.new_notebook()
+    original_source = "# sqlnbfmt: skip\n%%sql\nselect * from users where active = 1"
+    nb.cells = [nbf.new_code_cell(original_source)]
+    nb_path = tmp_path / "skip_magic.ipynb"
+    nbformat.write(nb, nb_path)
+
+    config = load_config()
+    changed = process_notebook(nb_path, config, "mysql", logger)
+    assert not changed
+
+    result_nb = nbformat.read(nb_path, as_version=4)
+    assert result_nb.cells[0].source == original_source
+
+
+def test_check_diff_combined(tmp_path, logger):
+    """--check --diff combined: no duplicate files, diff is printed."""
+    nb = nbf.new_notebook()
+    nb.cells = [
+        nbf.new_code_cell('query = "select id, name from users where active = 1"')
+    ]
+    nb_path = tmp_path / "combined.ipynb"
+    nbformat.write(nb, nb_path)
+
+    config = load_config()
+    # Simulate the combined logic from main()
+    diff_output = diff_notebook(nb_path, config, "mysql", logger)
+    assert diff_output  # has diff
+    # In the refactored main(), diff_notebook is used as the single source
+    # of truth for both --check and --diff, so no duplication possible.
+
+
+def test_main_diff_only_exits_zero(tmp_path):
+    """--diff alone always exits 0 even when changes are needed."""
+    import subprocess
+
+    nb = nbf.new_notebook()
+    nb.cells = [
+        nbf.new_code_cell('query = "select id, name from users where active = 1"')
+    ]
+    nb_path = tmp_path / "diff_exit.ipynb"
+    nbformat.write(nb, nb_path)
+
+    result = subprocess.run(
+        ["sqlnbfmt", "--diff", str(nb_path)],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0
+    assert "---" in result.stdout  # diff was printed
+
+
+def test_main_check_exits_one(tmp_path):
+    """--check exits 1 when formatting is needed."""
+    import subprocess
+
+    nb = nbf.new_notebook()
+    nb.cells = [
+        nbf.new_code_cell('query = "select id, name from users where active = 1"')
+    ]
+    nb_path = tmp_path / "check_exit.ipynb"
+    nbformat.write(nb, nb_path)
+
+    result = subprocess.run(
+        ["sqlnbfmt", "--check", str(nb_path)],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 1
+
+
+def test_main_check_diff_exits_one(tmp_path):
+    """--check --diff combined exits 1 and prints diff."""
+    import subprocess
+
+    nb = nbf.new_notebook()
+    nb.cells = [
+        nbf.new_code_cell('query = "select id, name from users where active = 1"')
+    ]
+    nb_path = tmp_path / "check_diff_exit.ipynb"
+    nbformat.write(nb, nb_path)
+
+    result = subprocess.run(
+        ["sqlnbfmt", "--check", "--diff", str(nb_path)],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 1
+    assert "---" in result.stdout  # diff was printed