sqlnbfmt 0.1.6__tar.gz → 0.1.8__tar.gz

{sqlnbfmt-0.1.6/sqlnbfmt.egg-info → sqlnbfmt-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlnbfmt
-Version: 0.1.6
+Version: 0.1.8
 Summary: A tool to format SQL code in Jupyter notebooks.
 Author-email: Qing <qingye779@gmail.com>
 License: MIT
@@ -10,6 +10,7 @@ Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: astor>=0.8.1
+Requires-Dist: black>=24.10.0
 Requires-Dist: nbformat>=5.10.4
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: sqlglot>=25.27.0

{sqlnbfmt-0.1.6 → sqlnbfmt-0.1.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlnbfmt"
-version = "0.1.6"
+version = "0.1.8"
 description = "A tool to format SQL code in Jupyter notebooks."
 authors = [
   { name = "Qing", email = "qingye779@gmail.com" }
@@ -14,6 +14,7 @@ classifiers = [
 ]
 dependencies = [
     "astor>=0.8.1",
+    "black>=24.10.0",
     "nbformat>=5.10.4",
     "pyyaml>=6.0.2",
     "sqlglot>=25.27.0",
@@ -38,4 +39,4 @@ markers = [
 ]
 
 [tool.setuptools.package-data]
-sqlnbfmt = [".pre-commit-hooks.yaml"]
+sqlnbfmt = [".pre-commit-hooks.yaml"]

{sqlnbfmt-0.1.6 → sqlnbfmt-0.1.8}/sqlnbfmt/formatter.py

@@ -7,6 +7,7 @@ from pathlib import Path
 from typing import Set, Optional, Dict, Union, Tuple
 
 import astor
+import black
 import nbformat
 import yaml
 from sqlglot import parse_one, errors
@@ -254,7 +255,7 @@ class SQLStringFormatter(ast.NodeTransformer):
                         formatted_sql = '\n'.join(indented_lines)
                         
                         # Reconstruct the string with appropriate quotes
-                        formatted_str = f'"""\n{formatted_sql}\n"""'
+                        formatted_str = f'"""\n{formatted_sql}\n{" " * 4}"""'
                     else:
                         # Ensure newlines are added around the formatted SQL
                         formatted_str = f'"""\n{formatted_sql}\n"""'
@@ -328,11 +329,18 @@ class SQLStringFormatter(ast.NodeTransformer):
 
                     new_node = ast.JoinedStr(values=new_values)
                     if '\n' in formatted_sql:
-                        formatted_fstring = ast.JoinedStr(values=[
-                            ast.Constant(value='\n'),
-                            *new_node.values,
-                            ast.Constant(value='\n')
-                        ])
+                        if in_function:
+                            formatted_fstring = ast.JoinedStr(values=[
+                                ast.Constant(value='\n'),
+                                *new_node.values,
+                                ast.Constant(value=f'\n{" " * 4}')
+                            ])
+                        else:
+                            formatted_fstring = ast.JoinedStr(values=[
+                                ast.Constant(value='\n'),
+                                *new_node.values,
+                                ast.Constant(value='\n')
+                            ])
                     else:
                         formatted_fstring = new_node
                     return formatted_fstring
@@ -477,6 +485,9 @@ def process_notebook(
                 if formatter.changed:
                     # Use astor to convert AST back to source code
                     formatted_code = astor.to_source(new_tree)
+                    # Now format with black
+                    formatted_code = black.format_str(formatted_code, mode=black.FileMode())
+                    
                     if formatted_code != original_code:
                         cell.source = formatted_code
                         changed = True

{sqlnbfmt-0.1.6 → sqlnbfmt-0.1.8/sqlnbfmt.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlnbfmt
-Version: 0.1.6
+Version: 0.1.8
 Summary: A tool to format SQL code in Jupyter notebooks.
 Author-email: Qing <qingye779@gmail.com>
 License: MIT
@@ -10,6 +10,7 @@ Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: astor>=0.8.1
+Requires-Dist: black>=24.10.0
 Requires-Dist: nbformat>=5.10.4
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: sqlglot>=25.27.0

{sqlnbfmt-0.1.6 → sqlnbfmt-0.1.8}/sqlnbfmt.egg-info/requires.txt

@@ -1,4 +1,5 @@
 astor>=0.8.1
+black>=24.10.0
 nbformat>=5.10.4
 pyyaml>=6.0.2
 sqlglot>=25.27.0

{sqlnbfmt-0.1.6 → sqlnbfmt-0.1.8}/tests/test_formatter.py

@@ -45,7 +45,8 @@ WHERE
     # Test Case 2: SELECT Query with F-String and Variable
     (
         [nbf.new_code_cell('''table_name = "users"\nquery = f"select * from {table_name}"''')],
-        [nbf.new_code_cell('''table_name = 'users'\nquery = f"""
+        [nbf.new_code_cell('''table_name = "users"
+query = f"""
 SELECT
   *
 FROM {table_name}
@@ -54,7 +55,8 @@ FROM {table_name}
     # Test Case 3: Multi-line SELECT Query with Variables
     (
         [nbf.new_code_cell('''columns = "id, name"\nquery = f"select {columns} from users where active = 1"''')],
-        [nbf.new_code_cell('''columns = 'id, name'\nquery = f"""
+        [nbf.new_code_cell('''columns = "id, name"
+query = f"""
 SELECT
   {columns}
 FROM users
@@ -70,18 +72,26 @@ WHERE
     # Test Case 5: SQL Query with Cell Magic Command
     (
         [nbf.new_code_cell('''%%sql\nselect * from users where active = 1''')],
-        [nbf.new_code_cell('''%%sql\nSELECT\n  *\nFROM users\nWHERE\n  active = 1''')],
+        [nbf.new_code_cell('''%%sql
+SELECT
+  *
+FROM users
+WHERE
+  active = 1''')],
     ),
-    # Test Case 6: INSERT Query with Variables (Updated to handle formatter limitations)
+    # Test Case 6: INSERT Query with Variables (Remains Unchanged)
     (
         [nbf.new_code_cell('''table_name = "users"\ncolumns = "(id, name)"\nvalues = "(1, 'Alice')"\nquery = f"insert into {table_name} {columns} values {values}"''')],
-        # Since the formatter cannot handle placeholders here, we expect the original code to remain unchanged
-        [nbf.new_code_cell('''table_name = "users"\ncolumns = "(id, name)"\nvalues = "(1, 'Alice')"\nquery = f"insert into {table_name} {columns} values {values}"''')],
+        [nbf.new_code_cell('''table_name = "users"
+columns = "(id, name)"
+values = "(1, 'Alice')"
+query = f"insert into {table_name} {columns} values {values}"''')],
     ),
     # Test Case 7: UPDATE Query with Variables
     (
         [nbf.new_code_cell('''table_name = "users"\nquery = f"update {table_name} set active = 0 where id = 1"''')],
-        [nbf.new_code_cell('''table_name = 'users'\nquery = f"""
+        [nbf.new_code_cell('''table_name = "users"
+query = f"""
 UPDATE {table_name} SET active = 0
 WHERE
   id = 1
@@ -136,6 +146,32 @@ SELECT
 FROM users
 """''')],
     ),
+    # Test Case 13: In-function SQL Query
+    (
+        [nbf.new_code_cell('''pd.read_sql("""select id, name from users where active = 1""")''')],
+        [nbf.new_code_cell('''pd.read_sql(
+    """
+    SELECT
+      id,
+      name
+    FROM users
+    WHERE
+      active = 1
+    """
+)''')],
+    ),
+    # Test Case 14: In-function SQL Query with F-String
+    (
+        [nbf.new_code_cell('''table_name = "users"\npd.read_sql(f"select * from {table_name}")''')],
+        [nbf.new_code_cell('''table_name = "users"
+pd.read_sql(
+    f"""
+    SELECT
+      *
+    FROM {table_name}
+    """
+)''')],
+    ),
 ])
 def test_sql_formatter(temp_nb_path, logger, input_cells, expected_cells):
     # Create a new notebook