diff --git a/zero_true/__init__.py b/zero_true/__init__.py index 75c2deb1..431fbe64 100644 --- a/zero_true/__init__.py +++ b/zero_true/__init__.py @@ -21,4 +21,14 @@ from zt_backend.models.components.html import HTML, pygwalker from zt_backend.models.state.state import state from zt_backend.models.state.state import global_state -from zt_backend.models.components.chat_ui import chat_ui \ No newline at end of file +from zt_backend.models.components.chat_ui import chat_ui +from zt_backend.utils.pyfile_parser import cell,notebook + +def sql(var_name,sql_string): + return(var_name,sql_string) + +def md(md_string): + return(md_string) + +def text(txt_string): + return(txt_string) diff --git a/zt_backend/main.py b/zt_backend/main.py index cbbb8e61..b6e4686b 100644 --- a/zt_backend/main.py +++ b/zt_backend/main.py @@ -2,7 +2,7 @@ from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from zt_backend.config import settings -from zt_backend.utils.notebook import get_notebook, write_notebook +from zt_backend.utils.notebook import get_notebook, write_notebook_to_python, write_notebook from zt_backend.utils.dependencies import parse_dependencies, write_dependencies from copilot.copilot import copilot_app import zt_backend.router as router @@ -45,9 +45,11 @@ def open_project(): except Exception as e: logger.error("Unexpected error with matplotlib configuration: %s", e) notebook_path = Path(settings.zt_path) / "notebook.ztnb" - if not notebook_path.exists(): + notebook_path = Path(settings.zt_path) / "notebook.py" + ztnb_path= Path(settings.zt_path) / "notebook.ztnb" + if not notebook_path.exists() and not ztnb_path.exists(): logger.info("No notebook file found, creating with empty notebook") - write_notebook() + write_notebook_to_python() requirements_path = Path(settings.zt_path) / "requirements.txt" if not requirements_path.exists(): logger.info("No requirements file found, creating empty file") diff --git a/zt_backend/models/components/card.py b/zt_backend/models/components/card.py index cb7ea81b..170833a2 100644 --- a/zt_backend/models/components/card.py +++ b/zt_backend/models/components/card.py @@ -25,4 +25,4 @@ class Card(ZTComponent): description="Density of the component", ) width: Optional[Union[int, str]] = Field("100%", description="Width of the card") - style: Optional[str] = Field("", description="CSS style to apply to the component") + style: Optional[str] = Field("", description="CSS style to apply to the card") diff --git a/zt_backend/models/state/notebook_state.py b/zt_backend/models/state/notebook_state.py index 59544652..085e420d 100644 --- a/zt_backend/models/state/notebook_state.py +++ b/zt_backend/models/state/notebook_state.py @@ -14,7 +14,7 @@ def __init__(self): self.notebook_db_path = notebook_db_path codeCell = notebook.CodeCell( - id=str(uuid.uuid4()), + id=str(uuid.uuid4()).split('-')[0], code="", components=[], variable_name="", diff --git a/zt_backend/router.py b/zt_backend/router.py index e55eee1c..c741f21f 100644 --- a/zt_backend/router.py +++ b/zt_backend/router.py @@ -54,7 +54,7 @@ from typing import Dict, Tuple, Optional import tempfile from zt_backend.utils.file_utils import * - +from zt_backend.utils.notebook import notebook_state router = APIRouter() manager = ConnectionManager() @@ -193,8 +193,9 @@ def run_all(): def create_cell(cellRequest: request.CreateRequest): if app_state.run_mode == "dev": logger.debug("Code cell addition request started") + num_cells = str(uuid.uuid4()).split('-')[0] createdCell = notebook.CodeCell( - id=str(uuid.uuid4()), + id="cell_"+str(num_cells), code="", components=[], output="", @@ -1063,7 +1064,7 @@ async def move_item(move_request: request.MoveItemRequest) -> Dict: ) # Protected files check - protected_files = ["requirements.txt", "notebook.ztnb", + protected_files = ["requirements.txt", "notebook.ztnb", "notebook.py", "zt_db.db", "zt_db.db.wal"] if source_path.name in protected_files: raise HTTPException( diff --git a/zt_backend/runner/execute_code.py b/zt_backend/runner/execute_code.py index b1c86435..a4b48600 100644 --- a/zt_backend/runner/execute_code.py +++ b/zt_backend/runner/execute_code.py @@ -21,6 +21,33 @@ now = datetime.now() logger = logging.getLogger("__name__") +def better_exec(code_, globals_=None, locals_=None, /, *, closure=None): + import ast + import linecache + + if not hasattr(better_exec, "saved_sources"): + old_getlines = linecache.getlines + better_exec.saved_sources = [] + + def patched_getlines(filename, module_globals=None): + if "")[0]) + return better_exec.saved_sources[index].splitlines(True) + else: + return old_getlines(filename, module_globals) + + linecache.getlines = patched_getlines + better_exec.saved_sources.append(code_) + exec( + compile( + ast.parse(code_), + filename=f"", + mode="exec", + ), + globals_ + ) + + initial_cell = request.CodeRequest( id="initial_cell", code=""" @@ -236,7 +263,7 @@ def flush(self): component_values=execution_state.component_values, ) execution_state.context_globals["exec_mode"] = True - exec(code_cell.code, temp_globals) + better_exec(code_cell.code,temp_globals) except Exception: logger.debug("Error during code execution") diff --git a/zt_backend/tests/test_convert.py b/zt_backend/tests/test_convert.py index 59162886..a1626534 100644 --- a/zt_backend/tests/test_convert.py +++ b/zt_backend/tests/test_convert.py @@ -1,51 +1,36 @@ from zt_backend.models.state.notebook_state import NotebookState -from zt_backend.models import notebook -import rtoml import subprocess from pathlib import Path import os - -notebook_state = NotebookState() +import importlib.util IPYNB_PATH = Path("zt_backend/tests/test_file.ipynb").resolve() -OUTPUT_PATH = Path("zt_backend/tests/notebook.ztnb").resolve() -NOTEBOOK_PATH = Path("zt_backend/tests/test_notebook.ztnb").resolve() +OUTPUT_PATH = Path("zt_backend/tests/notebook.py").resolve() +NOTEBOOK_PATH = Path("zt_backend/tests/test_notebook.py").resolve() + +def dynamic_import(module_path): + spec = importlib.util.spec_from_file_location("notebook_module", module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module def test_ipynb_to_ztnb(): + # Step 1: Convert the IPYNB file to a Python file convert = subprocess.Popen( ["zero-true", "jupyter-convert", IPYNB_PATH, OUTPUT_PATH] ) convert.wait() - with open(NOTEBOOK_PATH, "r", encoding="utf-8") as file: - expected_data = rtoml.loads(file.read().replace("\\", "\\\\")) - - expected_notebook_data = { - "notebookId": "test_id", - "notebookName": expected_data.get("notebookName", "Zero True"), - "userId": "", - "cells": { - f"{index}": notebook.CodeCell(id=f"{index}", **cell_data, output="") - for index, (cell_id, cell_data) in enumerate(expected_data["cells"].items()) - }, - } - expected_notebook = notebook.Notebook(**expected_notebook_data) - - with open(OUTPUT_PATH, "r", encoding="utf-8") as file: - output_data = rtoml.loads(file.read().replace("\\", "\\\\")) - - output_notebook_data = { - "notebookId": "test_id", - "notebookName": output_data.get("notebookName", "Zero True"), - "userId": "", - "cells": { - f"{index}": notebook.CodeCell(id=f"{index}", **cell_data, output="") - for index, (cell_id, cell_data) in enumerate(output_data["cells"].items()) - }, - } - output_notebook = notebook.Notebook(**output_notebook_data) - - assert expected_notebook == output_notebook - - os.remove(OUTPUT_PATH) \ No newline at end of file + # Step 2: Import the expected notebook from test_notebook.py + from test_notebook import notebook as expected_notebook + output_module = dynamic_import(OUTPUT_PATH) + output_notebook = output_module.notebook + # Step 3: Import the generated notebook from notebook.py + expected_notebook.notebookId='0' + output_notebook.notebookId='0' + # Step 4: Assert that the generated notebook matches the expected notebook + assert output_notebook == expected_notebook + + # Step 5: Clean up the generated file + os.remove(OUTPUT_PATH) diff --git a/zt_backend/tests/test_e2e.py b/zt_backend/tests/test_e2e.py index 41397780..8dad10db 100644 --- a/zt_backend/tests/test_e2e.py +++ b/zt_backend/tests/test_e2e.py @@ -1,5 +1,6 @@ import pytest import subprocess +import textwrap import os from selenium import webdriver from selenium.webdriver.common.by import By @@ -14,7 +15,6 @@ notebook_id = str(uuid.uuid4()) - expected_code = """ import zero_true as zt import time @@ -23,23 +23,24 @@ zt.TextInput(id='text')""" -notebook_str = '''notebookId = "''' + notebook_id + '''" -notebookName = "Zero True" +# Define the expected Python code for the notebook +notebook_str = f""" +import zero_true as zt +import time -[cells.57fbbd59-8f30-415c-87bf-8caae0374070] -cellName = "" -cellType = "code" -hideCell = "False" -hideCode = "False" -expandCode = "False" -showTable = "False" -nonReactive = "False" -code = """ -'''+expected_code+'''""" +def cell_0(): +"""+textwrap.indent(expected_code,' ')+""" -''' +notebook = zt.notebook( + id="{notebook_id}", + name="Zero True", + cells=[ + zt.cell(cell_0, type="code") + ] +) +""" -notebook_filename = "notebook.ztnb" +notebook_filename = "notebook.py" @pytest.fixture(scope="session", autouse=True) def start_stop_app(): @@ -81,6 +82,7 @@ def find_code_cells(driver): EC.presence_of_element_located((By.XPATH, "//div[contains(@id, 'codeCard')]")) ) code_cells = driver.find_elements(By.XPATH, "//div[contains(@id, 'codeCard')]") + print(code_cells) return code_cells def find_el_by_id_w_exception(driver,element_id): @@ -168,7 +170,7 @@ def test_notebook_loading(driver): def test_initial_code_cell(driver): code_cells = find_code_cells(driver) - assert len(code_cells) == 1 and code_cells[0].get_attribute('id') == 'codeCard57fbbd59-8f30-415c-87bf-8caae0374070', "Expected code cell not found." + assert len(code_cells) == 1 and code_cells[0].get_attribute('id') == 'codeCardcell_0', "Expected code cell not found." def test_intial_code_cell_content(driver): code_cells = find_code_cells(driver) @@ -226,7 +228,7 @@ def test_adding_new_code_cell(driver): code_cells = find_code_cells(driver) assert len(code_cells) == 2, "New code cell was not added" - assert code_cells[0].get_attribute('id') == 'codeCard57fbbd59-8f30-415c-87bf-8caae0374070', "Expected code cell not found" + assert code_cells[0].get_attribute('id') == 'codeCardcell_0', "Expected code cell not found" def test_execution_of_new_code_cell(driver): code_cells = find_code_cells(driver) @@ -339,7 +341,7 @@ def test_app_mode(driver): #assert that there is only cell in app mode because the second cell was hidden - cell_id_0 = '57fbbd59-8f30-415c-87bf-8caae0374070' + cell_id_0 = 'cell_0' assert len(code_cells) == 1 and code_cells[0].get_attribute('id') == f'codeCard{cell_id_0}', "Expected code cell not found." @@ -385,7 +387,7 @@ def test_deletion_of_new_code_cell(driver): delete_btn.click() time.sleep(2) code_cells = find_code_cells(driver) - assert len(code_cells) == 1 and code_cells[0].get_attribute('id') == 'codeCard57fbbd59-8f30-415c-87bf-8caae0374070', "Expected code cell not found." + assert len(code_cells) == 1 and code_cells[0].get_attribute('id') == 'codeCardcell_0', "Expected code cell not found." # test hiding code cell diff --git a/zt_backend/tests/test_notebook.py b/zt_backend/tests/test_notebook.py new file mode 100644 index 00000000..0172e61b --- /dev/null +++ b/zt_backend/tests/test_notebook.py @@ -0,0 +1,16 @@ +import zero_true as zt + +def cell_0(): + print('Hello, world!') + +def cell_1(): + zt.markdown("""# This is a markdown cell""") + +notebook = zt.notebook( + id='ca156d73-1c20-48c6-afbb-ae177c6bafa5', + name='Zero True', + cells=[ + zt.cell(cell_0, type='code'), + zt.cell(cell_1, type='markdown'), + ] +) diff --git a/zt_backend/tests/test_notebook.ztnb b/zt_backend/tests/test_notebook.ztnb deleted file mode 100644 index fadcb825..00000000 --- a/zt_backend/tests/test_notebook.ztnb +++ /dev/null @@ -1,35 +0,0 @@ -notebookId = "d8f7d425-bd7b-40ba-86be-f7911bd0643e" -notebookName = "Zero True" - -[cells.21623c58-7e2b-45ed-a440-f58f62b7288e] -cellName = "" -cellType = "code" -hideCell = "False" -hideCode = "False" -expandCode = "False" -showTable = "False" -nonReactive = "False" -code = """ -import zero_true as zt""" - -[cells.38d10a6f-3723-41df-a518-3887470cb78d] -cellName = "" -cellType = "code" -hideCell = "False" -hideCode = "False" -expandCode = "False" -showTable = "False" -nonReactive = "False" -code = """ -print('Hello, world!')""" - -[cells.39055a62-f0c8-42c4-9f17-e531f90e3298] -cellName = "" -cellType = "markdown" -hideCell = "False" -hideCode = "False" -expandCode = "False" -showTable = "False" -nonReactive = "False" -code = """ -# This is a markdown cell""" \ No newline at end of file diff --git a/zt_backend/utils/completions.py b/zt_backend/utils/completions.py index ceddc632..878804df 100644 --- a/zt_backend/utils/completions.py +++ b/zt_backend/utils/completions.py @@ -12,7 +12,7 @@ async def get_code_completions(cell_id: str, code: str, line: int, column: int) await text_document_did_change({ "textDocument": { - "uri": "file:///notebook.ztnb", + "uri": "file:///notebook.py", "version": 1 }, "contentChanges": [{ diff --git a/zt_backend/utils/notebook.py b/zt_backend/utils/notebook.py index b0e2af99..71089181 100644 --- a/zt_backend/utils/notebook.py +++ b/zt_backend/utils/notebook.py @@ -5,6 +5,7 @@ from zt_backend.models import notebook from zt_backend.utils.debounce import debounce from zt_backend.config import settings +from zt_backend.utils.pyfile_parser import update_notebook_file,load_notebook_from_file from pathlib import Path import logging import duckdb @@ -15,10 +16,88 @@ import copy import rtoml import ast +import re +import importlib logger = logging.getLogger("__name__") notebook_state = NotebookState() +def write_notebook_to_python(notebook_state=notebook_state): + update_notebook_file('notebook.py',notebook_state.zt_notebook) + + +def get_notebook_python(): + # Load the Python notebook file dynamically + notebook_path = Path(settings.zt_path) / "notebook.py" + + if notebook_path.exists(): + return(load_notebook_from_file(notebook_path)) + + +def get_ztnb(): + logger.debug("Notebook id is empty") + # If it doesn't exist in the database, load it from the TOML file + logger.info(f'{settings.zt_path}/notebook.py') + notebook_path = Path(settings.zt_path) / "notebook.ztnb" + + with notebook_path.open("r", encoding="utf-8") as project_file: + toml_data = rtoml.loads(project_file.read().replace("\\", "\\\\")) + + try: + # get notebook from the database + notebook_state.zt_notebook = get_notebook_db(toml_data["notebookId"]) + logger.debug( + "Notebook retrieved from db with id %s", toml_data["notebookId"] + ) + notebook_state_init() + return + except Exception as e: + logger.debug( + "Error loading notebook with id %s from db: %s", + toml_data["notebookId"], + traceback.format_exc(), + ) + pass + #Convert TOML data to a Notebook object + + notebook_data = { + "notebookId": toml_data["notebookId"], + "notebookName": toml_data.get("notebookName", "Zero True"), + "userId": "", + "cells": { + cell_id: notebook.CodeCell( + id=cell_id, + **{ + cell_key: cell_value + for cell_key, cell_value in cell_data.items() + if cell_key != "comments" + }, + output="", + comments={ + comment_id: notebook.Comment( + id=comment_id, + **{ + comment_key: comment_value + for comment_key, comment_value in comment_data.items() + if comment_key != "replies" + }, + replies={ + reply_id: notebook.Comment(id=reply_id, **reply_data) + for reply_id, reply_data in comment_data.get( + "replies", {} + ).items() + }, + ) + for comment_id, comment_data in cell_data.get( + "comments", {} + ).items() + }, + ) + for cell_id, cell_data in toml_data["cells"].items() + }, + } + return(notebook_data) + def get_notebook(id=""): if id != "": @@ -34,80 +113,48 @@ def get_notebook(id=""): ) try: - logger.debug("Notebook id is empty") - # If it doesn't exist in the database, load it from the TOML file - logger.info(f'{settings.zt_path}/notebook.ztnb') - notebook_path = Path(settings.zt_path) / "notebook.ztnb" - with notebook_path.open("r", encoding="utf-8") as project_file: - toml_data = rtoml.loads(project_file.read().replace("\\", "\\\\")) + ztnb_path = Path(settings.zt_path) / "notebook.ztnb" + notebook_path = Path(settings.zt_path) / "notebook.py" - try: - # get notebook from the database - notebook_state.zt_notebook = get_notebook_db(toml_data["notebookId"]) - logger.debug( - "Notebook retrieved from db with id %s", toml_data["notebookId"] + # Step 1: Check for the .ztnb file + if ztnb_path.exists(): + logger.info(f".ztnb file detected: {ztnb_path}") + # Load the .ztnb file + notebook_data = get_ztnb() + + # Save it as a Python file + update_notebook_file(notebook_path, notebook_state.zt_notebook) + + # Delete the .ztnb file + try: + ztnb_path.unlink() + logger.info(f".ztnb file deleted: {ztnb_path}") + except Exception as e: + logger.error(f"Error deleting .ztnb file: {e}") + + # Step 2: Load the Python file if it exists + elif notebook_path.exists(): + notebook_data = get_notebook_python() + + # Step 3: Handle loaded notebook data + if notebook_data is not None: + notebook_state.zt_notebook = notebook.Notebook(**notebook_data.model_dump()) + new_notebook = notebook_state.zt_notebook.model_dump_json() + + conn = duckdb.connect(notebook_state.notebook_db_path) + create_table_query = f"CREATE TABLE IF NOT EXISTS '{notebook_state.zt_notebook.notebookId}' (id STRING PRIMARY KEY, notebook STRING)" + conn.execute(create_table_query) + insert_query = f"INSERT OR REPLACE INTO '{notebook_state.zt_notebook.notebookId}' (id, notebook) VALUES (?, ?)" + conn.execute( + insert_query, [notebook_state.zt_notebook.notebookId, new_notebook] ) - notebook_state_init() - return - except Exception as e: + conn.close() + logger.debug( - "Error loading notebook with id %s from db: %s", - toml_data["notebookId"], - traceback.format_exc(), + "Notebook with id %s loaded from toml/python and new db entry created", + notebook_state.zt_notebook.notebookId, ) - pass - # Convert TOML data to a Notebook object - notebook_data = { - "notebookId": toml_data["notebookId"], - "notebookName": toml_data.get("notebookName", "Zero True"), - "userId": "", - "cells": { - cell_id: notebook.CodeCell( - id=cell_id, - **{ - cell_key: cell_value - for cell_key, cell_value in cell_data.items() - if cell_key != "comments" - }, - output="", - comments={ - comment_id: notebook.Comment( - id=comment_id, - **{ - comment_key: comment_value - for comment_key, comment_value in comment_data.items() - if comment_key != "replies" - }, - replies={ - reply_id: notebook.Comment(id=reply_id, **reply_data) - for reply_id, reply_data in comment_data.get( - "replies", {} - ).items() - }, - ) - for comment_id, comment_data in cell_data.get( - "comments", {} - ).items() - }, - ) - for cell_id, cell_data in toml_data["cells"].items() - }, - } - notebook_state.zt_notebook = notebook.Notebook(**notebook_data) - new_notebook = notebook_state.zt_notebook.model_dump_json() - conn = duckdb.connect(notebook_state.notebook_db_path) - create_table_query = f"CREATE TABLE IF NOT EXISTS '{notebook_state.zt_notebook.notebookId}' (id STRING PRIMARY KEY, notebook STRING)" - conn.execute(create_table_query) - insert_query = f"INSERT OR REPLACE INTO '{notebook_state.zt_notebook.notebookId}' (id, notebook) VALUES (?, ?)" - conn.execute( - insert_query, [notebook_state.zt_notebook.notebookId, new_notebook] - ) - conn.close() - logger.debug( - "Notebook with id %s loaded from toml and new db entry created", - toml_data["notebookId"], - ) - notebook_state_init() + notebook_state_init() except Exception as e: logger.error( "Error when loading notebook, return empty notebook: %s", @@ -317,7 +364,7 @@ def save_notebook(): insert_query = f"INSERT OR REPLACE INTO '{notebook_state.zt_notebook.notebookId}' (id, notebook) VALUES (?, ?)" conn.execute(insert_query, [notebook_state.zt_notebook.notebookId, new_notebook]) conn.close() - write_notebook() + write_notebook_to_python(notebook_state) def write_notebook(): @@ -386,6 +433,9 @@ def write_notebook(): logger.debug("Toml saved for notebook %s", notebook_state.zt_notebook.notebookId) + + + async def save_worker(save_queue): while True: message = await save_queue.get() diff --git a/zt_backend/utils/pyfile_parser.py b/zt_backend/utils/pyfile_parser.py new file mode 100644 index 00000000..f71732e0 --- /dev/null +++ b/zt_backend/utils/pyfile_parser.py @@ -0,0 +1,535 @@ +import ast +import inspect +import textwrap +import os +import importlib.util +from collections import OrderedDict, defaultdict +from pathlib import Path +from uuid import uuid4 +from zt_backend.models.notebook import Notebook, CodeCell +import astroid +from collections import defaultdict +import re +from zt_backend.runner.code_cell_parser import ( + get_imports, + get_defined_names, + get_loaded_names, + get_loaded_modules, + get_functions, +) + +def parse_cell(func): + """ + Inspect the function to detect: + - The function name (the cell_id). + - Whether its *only* statement is one of zt.sql(...), zt.markdown(...), or zt.text(...). + - If so, store only that string argument in cell_obj.code. + - Otherwise, store the full code block in cell_obj.code, excluding the function_def and the final return statement. + - Includes nested function definitions. + """ + WRAPPER_TO_TYPE = {"sql": "sql", "markdown": "markdown", "text": "text"} + + source = inspect.getsource(func) + source = textwrap.dedent(source) + tree = ast.parse(source) + + func_def = None + for node in tree.body: + if isinstance(node, ast.FunctionDef): + func_def = node + break + + if not func_def: + raise ValueError(f"No function definition found in {func.__name__}") + + cell_id = func_def.name + cell_type = "code" # default + + def filter_return_statement(node): + """Recursively process the AST to exclude the final return statement.""" + if isinstance(node, ast.FunctionDef): + node.body = [filter_return_statement(subnode) for subnode in node.body] + if node.body and isinstance(node.body[-1], ast.Return): + node.body = node.body[:-1] # Remove the final return statement + return node + + # Filter out the final return statement in the main function + func_def = filter_return_statement(func_def) + + # If there's only one statement and it's a zt.(...), extract it + if len(func_def.body) == 1 and isinstance(func_def.body[0], ast.Expr): + expr_node = func_def.body[0].value + if ( + isinstance(expr_node, ast.Call) + and isinstance(expr_node.func, ast.Attribute) + and isinstance(expr_node.func.value, ast.Name) + and expr_node.func.value.id == "zt" + ): + wrapper_name = expr_node.func.attr + if wrapper_name in WRAPPER_TO_TYPE: + cell_type = WRAPPER_TO_TYPE[wrapper_name] + if expr_node.args and isinstance(expr_node.args[0], ast.Constant): + return cell_id, expr_node.args[0].value, cell_type + + # Otherwise, reassemble the function's body excluding the definition and the final return + code_lines = source.splitlines() + start_line = func_def.lineno # Start after the function definition line + end_line = func_def.end_lineno + body_lines = code_lines[start_line:end_line] + + # Remove the last return line if it exists + dedented_body = textwrap.dedent("\n".join(body_lines)).strip() + body_lines = dedented_body.splitlines() + if body_lines and body_lines[-1].strip().startswith("return"): + body_lines = body_lines[:-1] # Exclude the final return line + + final_code = "\n".join(body_lines).strip() + return cell_id, final_code, cell_type + + +def cell(func, hidden=False, type='code', hide_code=False, expand_code=False, + show_table=False, non_reactive=False, cell_name=None,variable_name=""): + """ + Create a CodeCell from a function object, plus metadata. + We rely on parse_cell(...) to see if the function calls zt.sql, zt.markdown, etc. + If it's a single wrapper call, we store only the string argument in .code. + Otherwise, we store the entire block of code. + """ + cell_id, raw_code, detected_type = parse_cell(func) + var_name = variable_name + # If parse_cell found e.g. 'sql', 'markdown', or 'text', override 'type' + final_type = detected_type if detected_type != 'code' else type + return CodeCell( + id=cell_id, + code=raw_code, # raw code or text + output="", + cellName=(cell_name or ""), + hideCell=hidden, + hideCode=hide_code, + expandCode=expand_code, + showTable=show_table, + nonReactive=non_reactive, + cellType=final_type, + variable_name = var_name + ) + + +def notebook(id, name="Zero True", nonreactive=False, cells=None): + """ + Create a structured Notebook object from a list of CodeCell or function references. + """ + if cells is None: + cells = [] + return Notebook( + userId='', + notebookId=id, + notebookName=name, + nonreactive=nonreactive, + cells=OrderedDict((c.id, c) for c in cells) + ) + + +def load_notebook_from_file(file_path, notebook_variable_name="notebook"): + """ + Load and execute a Python file, returning the named notebook variable. + """ + module_name = os.path.splitext(os.path.basename(file_path))[0] + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + if hasattr(module, notebook_variable_name): + return getattr(module, notebook_variable_name) + else: + raise AttributeError(f"{notebook_variable_name} not found in {file_path}") + +def get_top_level_imports(source_code): + """ + Parse the source code and collect all top-level imports, including aliases. + """ + top_level_imports = set() + try: + tree = ast.parse(source_code) + for node in tree.body: + if isinstance(node, ast.Import): + for alias in node.names: + if alias.asname: + top_level_imports.add(f"import {alias.name} as {alias.asname}") + else: + top_level_imports.add(f"import {alias.name}") + elif isinstance(node, ast.ImportFrom): + module = node.module + for alias in node.names: + if alias.asname: + top_level_imports.add(f"from {module} import {alias.name} as {alias.asname}") + else: + top_level_imports.add(f"from {module} import {alias.name}") + except Exception as e: + print(f"[Warning] Could not parse top-level imports: {e}") + return top_level_imports + +def parse_cell_calls_from_notebook_line(source_code: str): + """ + Parse the source code to extract all cell IDs from a `notebook = zt.notebook(...)` definition. + Handles multi-line definitions and nested formatting. + """ + cell_ids = [] + try: + tree = ast.parse(source_code) + for node in ast.walk(tree): + # Look for the assignment `notebook = zt.notebook(...)` + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call): + if ( + isinstance(node.value.func, ast.Attribute) + and node.value.func.value.id == "zt" + and node.value.func.attr == "notebook" + ): + # Find the `cells=[...]` argument + for kw in node.value.keywords: + if kw.arg == "cells" and isinstance(kw.value, ast.List): + for elem in kw.value.elts: + if ( + isinstance(elem, ast.Call) + and isinstance(elem.func, ast.Attribute) + and elem.func.value.id == "zt" + and elem.func.attr == "cell" + ): + # Extract the first argument of zt.cell(...) + if elem.args and isinstance(elem.args[0], ast.Name): + cell_ids.append(elem.args[0].id) + except Exception as e: + print(f"[Warning] Could not parse cell IDs: {e}") + return cell_ids + + +def update_notebook_file(filepath, notebook_obj): + """ + Update or create a Python file to match the given Notebook’s cell definitions + and *physically remove* the function definition for any cell that was deleted + from the notebook. + This version: + - Keeps exactly one blank line before each cell (if needed). + - Removes old functions if they've been removed from the notebook. + - If a cell is recognized as markdown/sql/text, it writes zt.markdown(...), zt.sql(...), etc. in the file, + but the notebook only loads the raw string (see parse_cell). + """ + import re + import_line = "import zero_true as zt\n" + + # Helper function to ensure at most one blank line + def maybe_add_blank_line(line_list): + if line_list and line_list[-1].strip(): # last line not empty + line_list.append("\n") + + # Helper to build code block for each cell (in the file). + def build_cell_code_block(fn_name, cell_obj, def_line): + """ + Build the function code block for a given notebook cell. + Ensures imports, nested function variables, and internal definitions are excluded. + """ + return_line = " return" + filtered_arguments = [] # Final function arguments + filtered_returns = [] # Final return variables + + if cell_obj.cellType == "code": + try: + module = astroid.parse(cell_obj.code) + + # Gather imports + all_imports = get_imports(module) + + # Gather top-level defined variables (exclude nested function variables) + defined_names = get_defined_names(module) + function_names, _ = get_functions(module) + defined_names += function_names + + # Gather used variables + loaded_names = get_loaded_names(module, defined_names) + loaded_names += get_loaded_modules(module, all_imports) + + # Exclude duplicates while preserving order + seen = set() + filtered_arguments = [ + name for name in loaded_names + if name not in all_imports and name not in defined_names and name not in seen and not seen.add(name) + ] + + filtered_returns = [ + name for name in defined_names + if name not in all_imports and name not in seen and not seen.add(name) + ] + except Exception as e: + print(f"[Warning] Could not parse code in cell '{fn_name}': {e}") + filtered_arguments = [] + filtered_returns = [] + + # Update function signature + def_line = f"def {fn_name}({', '.join(filtered_arguments)}):" + + # Update the return statement + if filtered_returns: + return_line = f" return({', '.join(filtered_returns)})" + + # Build the full function + lines = [def_line] + + if cell_obj.cellType in ["markdown", "sql", "text"]: + lines.append(f" zt.{cell_obj.cellType}(\"\"\"{cell_obj.code}\"\"\")") + else: + for raw in cell_obj.code.splitlines(): + lines.append(f" {raw}") + + lines.append(return_line) + return lines + + # 1) Read existing lines or init + try: + with open(filepath, 'r') as f: + original_lines = f.readlines() + except FileNotFoundError: + original_lines = [import_line, "\n"] + + original_source = "".join(original_lines) + + # 1a) Gather old_cell_ids by scanning lines for 'notebook=zt.notebook(..., cells=[zt.cell(...), ...])' + + + old_cell_ids = parse_cell_calls_from_notebook_line(original_source) + + + # 2) Parse the file’s AST => gather all top-level function definitions + try: + tree = ast.parse(original_source) + except SyntaxError: + # If the file is partially corrupt, parse only the import line + tree = ast.parse(import_line) + + func_defs = [] + for node in tree.body: + if isinstance(node, ast.FunctionDef): + start_line = node.lineno - 1 + end_line = node.end_lineno - 1 + fn_name = node.name + func_defs.append((start_line, end_line, fn_name)) + + func_defs.sort(key=lambda x: x[0]) + + fn_name_map = defaultdict(list) + for (start_i, end_i, fn_name) in func_defs: + fn_name_map[fn_name].append((start_i, end_i)) + + new_cell_ids = set(notebook_obj.cells.keys()) + final_defs = [] + + # 3) For each function name, check if it’s recognized as an old cell + # - If old cell but not in new => remove + # - If old cell & in new => keep last definition => rewrite + # - Otherwise => keep as random code + for fn_name, ranges in fn_name_map.items(): + ranges.sort(key=lambda x: x[0]) + if fn_name in old_cell_ids: + if fn_name not in new_cell_ids: + # user deleted this cell => remove all definitions + for (s_i, e_i) in ranges: + final_defs.append((s_i, e_i, "_REMOVE_", True)) + else: + # keep last => rewrite; earlier are duplicates + last_start, last_end = ranges[-1] + final_defs.append((last_start, last_end, fn_name, True)) + for (s_i, e_i) in ranges[:-1]: + final_defs.append((s_i, e_i, "_DUPLICATE_", True)) + else: + # not recognized as old cell => keep as non-cell code + for (s_i, e_i) in ranges: + final_defs.append((s_i, e_i, fn_name, False)) + + final_defs.sort(key=lambda x: x[0]) + + updated_lines = [] + current_idx = 0 + handled_cells = set() + + i = 0 + while i < len(final_defs): + (start_i, end_i, fn_name, is_cell) = final_defs[i] + + # copy lines up to start_i + while current_idx < start_i and current_idx < len(original_lines): + updated_lines.append(original_lines[current_idx]) + current_idx += 1 + + if fn_name == "_DUPLICATE_": + current_idx = end_i + 1 + i += 1 + continue + if fn_name == "_REMOVE_": + current_idx = end_i + 1 + i += 1 + continue + + if is_cell: + # This is a cell in both old and new => we rewrite + if fn_name in new_cell_ids: + cell_obj = notebook_obj.cells[fn_name] + handled_cells.add(fn_name) + def_line = original_lines[start_i].rstrip("\n") + + # Add at most one blank line before the new version + maybe_add_blank_line(updated_lines) + + new_block = build_cell_code_block(fn_name, cell_obj, def_line) + for ln in new_block: + updated_lines.append(ln + "\n") + # Move on + current_idx = end_i + 1 + i += 1 + else: + # Non-cell code => keep as is + while current_idx <= end_i and current_idx < len(original_lines): + updated_lines.append(original_lines[current_idx]) + current_idx += 1 + i += 1 + + # Handle new cells that are not already handled + for cid, cell_obj in notebook_obj.cells.items(): + if cid not in handled_cells: + maybe_add_blank_line(updated_lines) + def_line = f"def {cid}():" + new_block = build_cell_code_block(cid, cell_obj, def_line) + for ln in new_block: + updated_lines.append(ln + "\n") + + # copy leftover lines + while current_idx < len(original_lines): + updated_lines.append(original_lines[current_idx]) + current_idx += 1 + + # Remove old notebook lines + notebook_start_pattern = re.compile(r"notebook\s*=\s*zt\.notebook\(") + inside_notebook = False + final_buf = [] + for line in updated_lines: + # Detect the start of the notebook definition + if notebook_start_pattern.match(line.strip()): + inside_notebook = True + continue # Skip the line + # Detect the end of the notebook definition + if inside_notebook and line.strip() == ")": + inside_notebook = False + continue # Skip the closing parenthesis + # If not inside a notebook definition, retain the line + if not inside_notebook: + final_buf.append(line) + updated_lines = final_buf + + # Now add brand-new cells that didn't exist previously + all_imports = set() # Collect all imports across cells + # Parse imports from the cell's code + for cid, cobj in notebook_obj.cells.items(): + try: + module = ast.parse(cobj.code) + cell_imports = set() + for node in ast.walk(module): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.asname: + cell_imports.add(f"import {alias.name} as {alias.asname}") + else: + cell_imports.add(f"import {alias.name}") + elif isinstance(node, ast.ImportFrom): + for alias in node.names: + if alias.asname: + cell_imports.add(f"from {node.module} import {alias.name} as {alias.asname}") + else: + cell_imports.add(f"from {node.module} import {alias.name}") + all_imports.update(cell_imports) + except Exception as e: + print(f"[Warning] Could not parse imports in cell '{cid}': {e}") + + existing_imports = get_top_level_imports(original_source) + all_imports -= existing_imports + + # Build the final notebook=zt.notebook(...) line + def format_cell_call(fn_name, cobj): + """ + zt.cell(fn_name, type='markdown', hidden=True, etc.) + We only set arguments that differ from default. + """ + call = f"zt.cell({fn_name}" + # If cellType != 'code' + if cobj.cellType in ["markdown", "sql", "text"]: + call += f", type='{cobj.cellType}'" + if cobj.hideCell: + call += ", hidden=True" + if cobj.hideCode: + call += ", hide_code=True" + if cobj.expandCode: + call += ", expand_code=True" + if cobj.showTable: + call += ", show_table=True" + if cobj.nonReactive: + call += ", non_reactive=True" + if cobj.variable_name: + call += f", variable_name='{cobj.variable_name}'" + if cobj.cellName: + call += f", cell_name='{cobj.cellName}'" + call += ")" + return call + + # Build the notebook arguments + nb_args = [f"id={notebook_obj.notebookId!r}"] + if notebook_obj.notebookName != "Zero True": + nb_args.append(f"name={notebook_obj.notebookName!r}") + + # Format the cells list, one cell per line with proper indentation + cell_lines = [ + " " + format_cell_call(cid, cobj) + "," # Indented for cells list + for cid, cobj in notebook_obj.cells.items() + ] + + # Combine everything into a properly formatted multi-line notebook definition + notebook_lines = [ + "notebook = zt.notebook(", + " " + ",\n ".join(nb_args) + ",", # Add notebook args + " cells=[", + *cell_lines, # Add each cell line + " ]", + ")", + ] + + # Clean up trailing blank lines + while updated_lines and not updated_lines[-1].strip(): + updated_lines.pop() + + # Ensure at most one blank line before the final line + maybe_add_blank_line(updated_lines) + updated_lines.append("\n".join(notebook_lines)) + + # Make sure we end with a newline + if not updated_lines[-1].endswith("\n"): + updated_lines[-1] += "\n" + + + original_source = "".join(original_lines) + + # Gather top-level imports from the file + + # Consolidate and add missing imports + # Ensure import zero_true is always at the top + if "import zero_true as zt" not in existing_imports: + updated_lines.insert(0, "import zero_true as zt\n") + existing_imports.add("import zero_true as zt") + + # Add imports below zero_true + # Consolidate and add missing imports + if all_imports: + for imp in sorted(all_imports): + if imp not in existing_imports: + updated_lines.insert(1, imp + "\n") # Insert after `import zero_true as zt` + existing_imports.add(imp + "\n") # Track the added import + + + filepath_temp = str(filepath).replace('.py', str(uuid4()) + '.py') + with open(filepath_temp, 'w') as f: + f.writelines(updated_lines) + os.replace(filepath_temp, filepath) diff --git a/zt_cli/cli.py b/zt_cli/cli.py index d3ac7ac5..5ab9c2ec 100644 --- a/zt_cli/cli.py +++ b/zt_cli/cli.py @@ -229,84 +229,72 @@ def notebook( @cli_app.command() def jupyter_convert( ipynb_path: Annotated[str, typer.Argument(help="The path to the .ipynb file")], - ztnb_path: Annotated[ - Optional[str], typer.Argument(help="The path to the output .ztnb file") - ] = "notebook.ztnb", + nb_path: Annotated[ + Optional[str], typer.Argument(help="The path to the output .py file") + ] = "notebook.py", ): """ Convert a Jupyter notebook to a Zero-True notebook. """ # Add notebook.ztnb if not specified in the output path - if not ztnb_path.endswith("notebook.ztnb"): - ztnb_path = os.path.join(ztnb_path, "notebook.ztnb") + if not nb_path.endswith("notebook.py"): + nb_path = os.path.join(nb_path, "notebook.py") try: with open(ipynb_path, "r", encoding="utf-8") as f: - notebook = json.loads(f.read()) + notebook = json.load(f) except Exception as e: typer.echo(f"Error occured: {e}") return else: - output = [] + output = jupyter_convert_func(notebook) - output.append(f'notebookId = "{uuid.uuid4()}"') - output.append('notebookName = "Zero True"') - output.append("") - output.extend( - line for line in create_ztnb_cell('"code"', ["import zero_true as zt"]) - ) - - # Create only code or markdown cells - for cell in notebook["cells"]: - if cell["cell_type"] in ["code", "markdown"]: - output.extend( - line - for line in create_ztnb_cell( - f'"{cell["cell_type"]}"', cell["source"] - ) - ) - - with open(ztnb_path, "w", encoding="utf-8") as f: + with open(nb_path, "w", encoding="utf-8") as f: for item in output: f.write(item + "\n") - typer.echo(f"Successfully converted {ipynb_path} to {ztnb_path}") + typer.echo(f"Successfully converted {ipynb_path} to {nb_path}") return +def jupyter_convert_func(notebook): + output_lines = [] + + # Add notebook metadata + notebook_id = str(uuid.uuid4()) + output_lines.append(f"import zero_true as zt") + output_lines.append("") + + # Generate Python functions for each cell + for index, cell in enumerate(notebook["cells"], start=0): + cell_id = f"cell_{index}" + if cell["cell_type"] == "code": + output_lines.append(f"def {cell_id}():") + # Join the source lines with proper indentation + source_code = "".join(cell["source"]) + for line in source_code.splitlines(): + output_lines.append(f" {line}") + output_lines.append("") # Add an empty line after the function + + elif cell["cell_type"] == "markdown": + markdown_content = "".join(cell["source"]).strip().replace('"""', "'''") + output_lines.append(f"def {cell_id}():") + output_lines.append(f" zt.markdown(\"\"\"{markdown_content}\"\"\")") + output_lines.append("") # Add an empty line after the function + + # Add notebook definition + output_lines.append(f"notebook = zt.notebook(") + output_lines.append(f" id='{notebook_id}',") + output_lines.append(f" name='Zero True',") + output_lines.append(f" cells=[") + for index in range(0, len(notebook["cells"])): + output_lines.append(f" zt.cell(cell_{index}, type='{'markdown' if notebook['cells'][index]['cell_type'] == 'markdown' else 'code'}'),") + output_lines.append(f" ]") + output_lines.append(f")") + + return output_lines # Return the list of output lines directly -def create_ztnb_cell(cell_type, source): - - common_attributes = { - "cellName": '""', - "cellType": '"code"', - "hideCell": '"False"', - "hideCode": '"False"', - "expandCode": '"False"', - "showTable": '"False"', - "nonReactive": '"False"', - "code": '"""', - } - - cell_content = [] - cell_content.append(f"[cells.{uuid.uuid4()}]") - - for key, value in common_attributes.items(): - if key == "cellType": - cell_content.append(f"{key} = {cell_type}") - else: - cell_content.append(f"{key} = {value}") - - for line in source: - if cell_type == '"code"': - escaped_code = line.encode().decode("unicode_escape").replace('"""', "'''") - cell_content.append(escaped_code) - else: - cell_content.append(line) - cell_content[-1] = cell_content[-1] + '"""' - cell_content.append("") - return cell_content if __name__ == "__main__": diff --git a/zt_frontend/src/App.vue b/zt_frontend/src/App.vue index f26f4b77..e8f1041e 100644 --- a/zt_frontend/src/App.vue +++ b/zt_frontend/src/App.vue @@ -1031,7 +1031,7 @@ export default { { doc: { version: 1, - uri: "file:///notebook.ztnb", + uri: "file:///notebook.py", position: { line: this.concatenatedCodeCache.length + line, character: column, diff --git a/zt_frontend/src/components/FileExplorer.vue b/zt_frontend/src/components/FileExplorer.vue index cb6dbe67..99995088 100644 --- a/zt_frontend/src/components/FileExplorer.vue +++ b/zt_frontend/src/components/FileExplorer.vue @@ -116,7 +116,7 @@ export default defineComponent({ const searchResults = ref([]); const isSearching = ref(false); - const protectedFiles = ref(["requirements.txt", "notebook.ztnb","zt_db.db","zt_db.db.wal"]); + const protectedFiles = ref(["requirements.txt", "notebook.ztnb","notebook.py","zt_db.db","zt_db.db.wal"]); const isProtectedFile = (filename: string) => { return protectedFiles.value.includes(filename); diff --git a/zt_frontend/src/components/FileTreeNode.vue b/zt_frontend/src/components/FileTreeNode.vue index ea482c55..d905a08d 100644 --- a/zt_frontend/src/components/FileTreeNode.vue +++ b/zt_frontend/src/components/FileTreeNode.vue @@ -124,7 +124,7 @@ export default defineComponent({ const isDragging = ref(false); const isDragOver = ref(false); const isDragTarget = ref(false); - const protectedFiles = ref(["requirements.txt", "notebook.ztnb", "zt_db.db", "zt_db.db.wal"]); + const protectedFiles = ref(["requirements.txt", "notebook.ztnb","notebook.py", "zt_db.db", "zt_db.db.wal"]); watch(menuOpen, (newVal) => { if (newVal) { diff --git a/zt_frontend/src/plugins/vuetify.ts b/zt_frontend/src/plugins/vuetify.ts index a7c9e5b2..accd0785 100644 --- a/zt_frontend/src/plugins/vuetify.ts +++ b/zt_frontend/src/plugins/vuetify.ts @@ -57,6 +57,7 @@ const vuetify = createVuetify({ VCard: { color: "bluegrey-darken-4", class: "scroll" + }, VDivider: { class: 'border-opacity-100'