setta 0.0.1.dev0__py3-none-any.whl → 0.0.2.dev0__py3-none-any.whl

setta/tasks/fns/textFieldInitializeCode.py

@@ -0,0 +1,99 @@
+import json
+
+from pydantic import BaseModel
+
+from setta.code_gen.create_runnable_scripts import generate_final_code_for_sections
+from setta.code_gen.export_selected import export_selected
+from setta.code_gen.find_placeholders import tp
+from setta.tasks.task_runner import RunType
+from setta.utils.constants import C
+
+from .utils import (
+    TaskDefinition,
+    send_document_open_request_with_diagnostics_with_code_dict,
+)
+
+
+class TextFieldInitializeCodeContent(BaseModel):
+    project: dict
+    sectionId: str
+    paramInfoId: str | None
+
+
+async def _textFieldInitializeCode(message, lsp_writers):
+    message = TextFieldInitializeCodeContent.parse_obj(message.content)
+    project = message.project
+    exporter_obj = export_selected(
+        project, always_export_args_objs=True, force_include_template_var=True
+    )
+
+    var_name = exporter_obj.var_name_mapping[
+        (message.sectionId, message.paramInfoId, False)
+    ]
+    projectConfigName = project["projectConfig"]["name"]
+    code_dict = await generate_final_code_for_sections(
+        project,
+        exporter_obj,
+        lsp_writers=lsp_writers,
+        folder_path=lsp_writers["full"].project_code_folder(projectConfigName),
+        push_var_name=var_name,
+    )
+    send_document_open_request_with_diagnostics_with_code_dict(
+        lsp_writers["full"],
+        code_dict,
+        exporter_obj,
+        projectConfigName,
+    )
+
+    gen_code_section_details = exporter_obj.code_gen_template_var_section_details[
+        "section"
+    ]
+    gen_code_dict = code_dict[gen_code_section_details["id"]]
+    var_name_to_decl_position = gen_code_dict["var_name_to_decl_rel_position_dict"][
+        tp(C.SETTA_GENERATED_PYTHON)
+    ]
+    pushed_var_name_position = var_name_to_decl_position[var_name]
+    generated_var_names = list(var_name_to_decl_position.keys())
+    infoToFullName = {
+        tuple(json.loads(k)): v for k, v in project["infoToFullName"].items()
+    }
+    referencable_var_names = [
+        infoToFullName[exporter_obj.var_name_reverse_mapping[v]]
+        for v, (startPos, _, _) in var_name_to_decl_position.items()
+        if startPos < pushed_var_name_position[0]
+    ]
+
+    return {
+        "cache": [
+            gen_code_dict,
+            var_name,
+            referencable_var_names,
+            generated_var_names,
+            exporter_obj,
+            project["fullNameToInfo"],
+        ],
+    }
+
+
+def get_cache_key(projectConfigId, sectionId, paramInfoId):
+    return f"textFieldInitializeCode-{projectConfigId}-{sectionId}-{paramInfoId}"
+
+
+def write_cache(cache, message, result):
+    if result.get("cache") is None:
+        return
+    message = TextFieldInitializeCodeContent.parse_obj(message.content)
+    key = get_cache_key(
+        message.project["projectConfig"]["id"], message.sectionId, message.paramInfoId
+    )
+    cache[key] = result["cache"]
+
+
+textFieldInitializeCode = TaskDefinition(
+    name="textFieldInitializeCode",
+    return_message_type="textFieldInitializeCodeReturn",
+    fn=_textFieldInitializeCode,
+    write_cache=write_cache,
+    run_as=RunType.NONE,
+    server_dependencies=["lsp_writers"],
+)

setta/tasks/fns/typeCheck.py

@@ -0,0 +1,40 @@
+from setta.code_gen.create_runnable_scripts import generate_final_code_for_sections
+from setta.code_gen.export_selected import export_selected
+from setta.tasks.task_runner import RunType
+
+from .utils import (
+    TaskDefinition,
+    send_document_open_request_with_diagnostics_with_code_dict,
+)
+
+
+# this just initializes a document which gets the lsp to send back diagnostics
+async def _typeCheck(message, lsp_writers):
+    project = message.content
+    exporter_obj = export_selected(
+        project, always_export_args_objs=False, force_include_template_var=True
+    )
+    projectConfigName = project["projectConfig"]["name"]
+    code_dict = await generate_final_code_for_sections(
+        project,
+        exporter_obj,
+        lsp_writers=lsp_writers,
+        folder_path=lsp_writers["full"].project_code_folder(projectConfigName),
+    )
+    send_document_open_request_with_diagnostics_with_code_dict(
+        lsp_writers["full"],
+        code_dict,
+        exporter_obj,
+        projectConfigName,
+    )
+
+    return None
+
+
+typeCheck = TaskDefinition(
+    name="typeCheck",
+    return_message_type="typeCheckReturn",
+    fn=_typeCheck,
+    run_as=RunType.NONE,
+    server_dependencies=["lsp_writers"],
+)

setta/tasks/fns/utils.py

@@ -0,0 +1,134 @@
+import asyncio
+import random
+import string
+from typing import Any, Callable, List, Optional
+
+from pydantic import BaseModel
+
+from setta.tasks.task_runner import RunType
+from setta.utils.constants import C
+from setta.utils.utils import multireplace
+
+
+def send_document_open_request_with_diagnostics_with_code_dict(
+    lsp_writer,
+    code_dict,
+    exporter_obj,
+    projectConfigName,
+):
+    for codeSectionId, details in code_dict.items():
+        if details["codeLanguage"] != "python":
+            continue
+        asyncio.create_task(
+            send_document_open_request_with_diagnostics(
+                lsp_writer,
+                details["code"],
+                exporter_obj,
+                details["ref_template_var_positions"],
+                projectConfigName,
+                details["sanitized_full_name"],
+                codeSectionId,
+            )
+        )
+
+
+async def send_document_open_request_with_diagnostics(
+    lsp_writer,
+    code,
+    exporter_obj,
+    ref_template_var_positions,
+    projectConfigName,
+    sectionName,
+    codeSectionId,
+):
+    code_id = lsp_writer.construct_code_id(
+        projectConfigName,
+        sectionName,
+    )
+    await lsp_writer.send_document_open_request(
+        code_id=code_id,
+        code=code,
+        codeJson=exporter_obj.output if exporter_obj else None,
+        ref_template_var_positions=ref_template_var_positions,
+        needsDiagnostics=True,
+        projectConfigId=exporter_obj.p["projectConfig"]["id"],
+        codeSectionId=codeSectionId,
+    )
+
+    return code_id
+
+
+def prepare_code_for_scope_checking(fullText, templateVars, target_template_var):
+    the_special_var_name = None
+    for t in templateVars:
+        startPos = t["startPos"]
+        keyword = t["keyword"]
+        endPos = startPos + len(keyword)
+        if t is target_template_var:
+            # need to assign something to the replacement variable so that ast considers it a variable
+            decl = "=None"
+            # the [:-len(decl)] is to keep the full declaration the same length as the original keyword
+            the_special_var_name = generate_unique_variable(fullText, keyword)[
+                : -len(decl)
+            ]
+            replacement = the_special_var_name + decl
+        else:
+            # variable assignment not necessary for the keywords that we don't care about
+            replacement = generate_unique_variable(fullText, keyword)
+        fullText = fullText[:startPos] + replacement + fullText[endPos:]
+
+    return fullText, the_special_var_name
+
+
+# generate valid variable names because for example $SETTA_GENERATED_PYTHON is not valid
+def replace_template_vars_with_random_names(candidateTemplateVars, fullText):
+    var_name_mapping = {}
+    reverse_var_name_mapping = {}
+    keyword_to_pos = {}
+    for keyword, pos in candidateTemplateVars:
+        new_keyword = generate_unique_variable(fullText, keyword)
+        var_name_mapping[keyword] = new_keyword
+        reverse_var_name_mapping[new_keyword] = keyword
+        # we only need 1 position of the keyword, even if there are multiple positions
+        keyword_to_pos[keyword] = pos
+
+    code = multireplace(fullText, var_name_mapping)
+    return code, list(keyword_to_pos.values()), reverse_var_name_mapping
+
+
+def generate_unique_variable(code, keyword):
+    length = len(keyword)
+    chars = string.ascii_lowercase + string.digits + "_"
+
+    while True:
+        # Ensure the variable starts with a letter or underscore
+        new_var = random.choice(string.ascii_lowercase + "_")
+        # Fill the rest with random characters
+        new_var += "".join(random.choice(chars) for _ in range(length - 1))
+
+        if new_var.isidentifier() and new_var not in code:
+            return new_var
+
+
+def do_nothing(*args, **kwargs):
+    pass
+
+
+class TaskDefinition(BaseModel):
+    name: Optional[str] = ""
+    return_message_type: str
+    fn: Callable
+    read_cache: Optional[Callable] = do_nothing
+    write_cache: Optional[Callable] = do_nothing
+    run_as: Optional[RunType] = RunType.NONE
+    dependencies: Optional[List[str]] = None
+    server_dependencies: Optional[List[str]] = []
+
+
+class TaskMessage(BaseModel):
+    id: str
+    content: Any
+
+
+class SettaInMemoryFn(TaskDefinition):
+    return_message_type: str = C.WS_IN_MEMORY_FN_RETURN

setta/tasks/task_runner.py

@@ -0,0 +1,29 @@
+import asyncio
+import inspect
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+from enum import Enum
+
+
+class RunType(Enum):
+    SUBPROCESS = "SUBPROCESS"
+    THREAD = "THREAD"
+    NONE = "NONE"
+
+
+class TaskRunner:
+    def __init__(self):
+        self.thread_executor = ThreadPoolExecutor(max_workers=2)
+        self.process_executor = ProcessPoolExecutor(max_workers=2)
+
+    async def run(self, fn, fn_args, run_as):
+        if inspect.iscoroutinefunction(fn):
+            return await fn(*fn_args)
+
+        if run_as == RunType.NONE:
+            return fn(*fn_args)
+        elif run_as == RunType.THREAD:
+            executor = self.thread_executor
+        elif run_as == RunType.SUBPROCESS:
+            executor = self.process_executor
+
+        return await asyncio.get_running_loop().run_in_executor(executor, fn, *fn_args)

setta/tasks/tasks.py

@@ -0,0 +1,152 @@
+import asyncio
+import logging
+import queue
+import threading
+from typing import Dict
+
+from setta.database.utils import create_new_id
+
+from . import fns
+from .fns.utils import TaskDefinition, TaskMessage
+from .task_runner import RunType, TaskRunner
+from .utils import SettaInMemoryFnSubprocess, add_fns_from_module
+
+logger = logging.getLogger(__name__)
+
+
+class Tasks:
+    def __init__(self, lsp_writers):
+        self.lsp_writers = lsp_writers
+        self.task_runner = TaskRunner()
+        self.cache = {}
+        self.fns: Dict[str, TaskDefinition] = {}
+        self.in_memory_subprocess = SettaInMemoryFnSubprocess()
+        self.websockets = []  # Store the websocket connections
+        add_fns_from_module(self.fns, fns)
+
+        # Start stdout listener thread
+        self._stop_event = asyncio.Event()
+        self.stdout_queue = queue.Queue()  # regular Queue
+        self._stdout_processor_task = None
+        self.stdout_thread = threading.Thread(target=self._stdout_listener, daemon=True)
+        self.stdout_thread.start()
+
+    # Backend Changes (Tasks class)
+    async def connect(self, websocket):
+        # Accept the new connection
+        await websocket.accept()
+        self.websockets.append(websocket)
+
+        # Start the processor task if it's not running
+        if self._stdout_processor_task is None or self._stdout_processor_task.done():
+            self._stdout_processor_task = asyncio.create_task(
+                self._process_stdout_queue()
+            )
+
+    async def disconnect(self, websocket):
+        self.websockets.remove(websocket)
+        if len(self.websockets) == 0:
+            # Cancel the processor task
+            if self._stdout_processor_task and not self._stdout_processor_task.done():
+                self._stdout_processor_task.cancel()
+                try:
+                    await self._stdout_processor_task
+                except asyncio.CancelledError:
+                    pass
+                self._stdout_processor_task = None
+
+    def _stdout_listener(self):
+        while not self._stop_event.is_set():
+            try:
+                stdout_data = self.in_memory_subprocess.stdout_parent_conn.recv()
+                self.stdout_queue.put(stdout_data)  # simple put, no async needed
+            except Exception as e:
+                if self._stop_event.is_set():
+                    break
+                logger.debug(f"Error in stdout listener: {e}")
+
+    async def _process_stdout_queue(self):
+        while not self._stop_event.is_set():
+            try:
+                if self._stop_event.is_set():
+                    break
+                if len(self.websockets) > 0:
+                    stdout_data = self.stdout_queue.get_nowait()
+                    stdout_data = stdout_data.replace("\n", "\r\n")
+                    for w in self.websockets:
+                        await w.send_text(stdout_data)
+                    self.stdout_queue.task_done()
+            except queue.Empty:
+                await asyncio.sleep(0.1)  # Check for connection every 100ms
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                if self._stop_event.is_set():
+                    break
+                logger.debug(f"Error processing stdout: {e}")
+
+    async def __call__(self, fn_name, message: TaskMessage):
+        if fn_name in self.fns:
+            return await self.call_regular_fn(fn_name, message)
+        return await self.call_in_memory_subprocess_fn(fn_name, message)
+
+    async def call_regular_fn(self, fn_name, message: TaskMessage):
+        fn = self.fns[fn_name]
+        fn.read_cache(self.cache, message)
+        fn_args = [message, *(getattr(self, k) for k in fn.server_dependencies)]
+        result = await self.task_runner.run(fn.fn, fn_args, fn.run_as)
+        fn.write_cache(self.cache, message, result)
+        if result is None:
+            result = {}
+        result["messageType"] = fn.return_message_type
+        return result
+
+    async def call_in_memory_subprocess_fn(self, fn_name, message: TaskMessage):
+        self.in_memory_subprocess.parent_conn.send(
+            {"type": "call", "fn_name": fn_name, "message": message}
+        )
+        result = await self.task_runner.run(
+            self.in_memory_subprocess.parent_conn.recv, [], RunType.THREAD
+        )
+        # if result["status"] == "success":
+        # for x in result["content"]:
+        # x.setdefault("sectionType", default_section_type(x["type"]))
+        if result["status"] != "success":
+            result["content"] = {}
+
+        return {"content": result["content"], "messageType": result["messageType"]}
+
+    async def add_custom_fns(self, code_list, to_cache):
+        error_msgs = {}
+        task_metadata = {}
+        initial_content = []
+        for c in code_list:
+            # Send import request to subprocess
+            self.in_memory_subprocess.parent_conn.send(
+                {
+                    "type": "import",
+                    "code": c["code"],
+                    "module_name": c["module_name"],
+                    "to_cache": to_cache,
+                }
+            )
+            result = await self.task_runner.run(
+                self.in_memory_subprocess.parent_conn.recv, [], RunType.THREAD
+            )
+            if result["status"] == "success":
+                task_metadata.update(result["content"])
+            else:
+                error_msgs[c["module_name"]] = result["error"]
+
+        for k in task_metadata.keys():
+            task_output = await self(k, TaskMessage(id=create_new_id(), content={}))
+            initial_content.extend(task_output["content"])
+
+        return task_metadata, error_msgs, initial_content
+
+    def close(self):
+        self._stop_event.set()
+        self.in_memory_subprocess.close()
+        self.stdout_thread.join()
+        if self._stdout_processor_task:
+            self._stdout_processor_task.cancel()

setta/tasks/utils.py

@@ -0,0 +1,178 @@
+import importlib.util
+import json
+import multiprocessing
+import os
+import sys
+import traceback
+import uuid
+
+from setta.tasks.fns.utils import TaskDefinition
+from setta.utils.utils import nested_access
+
+
+def import_code_from_string(code_string, module_name=None, add_to_sys_modules=True):
+    # Generate a unique module name if one isn't provided
+    if module_name is None:
+        module_name = f"setta_dynamic_module_{uuid.uuid4().hex}"
+
+    # Add current directory to sys.path if it's not already there
+    current_dir = os.getcwd()
+    if current_dir not in sys.path:
+        sys.path.insert(0, current_dir)
+
+    spec = importlib.util.spec_from_loader(module_name, loader=None)
+
+    # Create a new module based on the spec
+    module = importlib.util.module_from_spec(spec)
+
+    # Optionally add the module to sys.modules
+    if add_to_sys_modules:
+        print(f"adding {module_name} to sys.modules", flush=True)
+        sys.modules[module_name] = module
+
+    # Compile the code string
+    code_object = compile(code_string, module_name, "exec")
+
+    # Execute the compiled code object in the module's namespace
+    exec(code_object, module.__dict__)
+
+    return module
+
+
+class SettaInMemoryFnSubprocess:
+    def __init__(self):
+        self.parent_conn, self.child_conn = multiprocessing.Pipe()
+        self.process = multiprocessing.Process(target=self._subprocess_main)
+        self.stdout_parent_conn, self.stdout_child_conn = multiprocessing.Pipe()
+        self.process.daemon = True  # Ensure process dies with parent
+        self.process.start()
+
+    def _subprocess_main(self):
+        """Main loop in subprocess that handles all requests"""
+        # Initialize store for imported modules
+        fns_dict = {}
+        cache = {}
+
+        class OutputCapture:
+            def __init__(self, stdout_pipe):
+                self.stdout_pipe = stdout_pipe
+
+            def write(self, text):
+                self.stdout_pipe.send(text)
+
+            def flush(self):
+                pass
+
+        # Redirect stdout as soon as subprocess starts
+        output_capture = OutputCapture(self.stdout_child_conn)
+        sys.stdout = output_capture
+        sys.stderr = output_capture
+
+        while True:
+            msg = self.child_conn.recv()  # Wait for requests
+            msg_type = msg["type"]
+            return_message_type = None
+
+            if msg_type == "shutdown":
+                break
+
+            try:
+                if msg_type == "import":
+                    code = msg["code"]
+                    module_name = msg["module_name"]
+                    # Import and store module
+                    module = import_code_from_string(code, module_name)
+                    added_fn_names = add_fns_from_module(fns_dict, module, module_name)
+                    task_metadata = {}
+                    for k in added_fn_names:
+                        cache[k] = msg["to_cache"]
+                        task_metadata[k] = get_task_metadata(fns_dict[k], cache[k])
+
+                    self.child_conn.send(
+                        {
+                            "status": "success",
+                            "content": task_metadata,
+                        }
+                    )
+
+                elif msg_type == "call":
+                    fn_name = msg["fn_name"]
+                    message = self.process_message(fn_name, msg["message"], cache)
+                    fn = fns_dict[fn_name]
+                    result = fn.fn(message)
+                    return_message_type = fn.return_message_type
+
+                    self.child_conn.send(
+                        {
+                            "status": "success",
+                            "content": result,
+                            "messageType": return_message_type,
+                        }
+                    )
+
+            except Exception as e:
+                traceback.print_exc()
+                self.child_conn.send(
+                    {
+                        "status": "error",
+                        "error": str(e),
+                        "messageType": return_message_type,
+                    }
+                )
+
+    def close(self):
+        try:
+            self.parent_conn.send({"type": "shutdown"})
+            self.process.join(timeout=1.0)
+        except:
+            pass
+
+        if self.process.is_alive():
+            self.process.terminate()
+            self.process.join()
+
+        # Close both sets of connections
+        self.parent_conn.close()
+        self.child_conn.close()
+        self.stdout_parent_conn.close()
+        self.stdout_child_conn.close()
+
+    def process_message(self, fn_name, message, cache):
+        if fn_name in cache:
+            exporter_obj = cache[fn_name]
+            for k, v in message.content.items():
+                nice_str = exporter_obj.var_name_mapping[tuple(json.loads(k))]
+                p_dict, key = nested_access(exporter_obj.output, nice_str)
+                p_dict[key] = v
+            message.content = exporter_obj.output
+        return message.content
+
+
+def add_fns_from_module(fns_dict, module, module_name=None):
+    count = 1
+    added_fn_names = []
+    for attr in vars(module).values():
+        if isinstance(attr, TaskDefinition):
+            if not attr.name:
+                if not module_name:
+                    raise ValueError(
+                        "Either TaskDefinition must be initialized with a name, or task_name must be specified"
+                    )
+                attr.name = f"{module_name}-{count}"
+                count += 1
+            fns_dict[attr.name] = attr
+            added_fn_names.append(attr.name)
+    return added_fn_names
+
+
+def get_task_metadata(in_memory_fn, exporter_obj):
+    # None means run the task on every change
+    if in_memory_fn.dependencies is None:
+        dependencies = None
+    # Empty array means only run when the task imported.
+    # Non-empty array means run when specified dependencies update.
+    else:
+        dependencies = [
+            exporter_obj.var_name_reverse_mapping[d] for d in in_memory_fn.dependencies
+        ]
+    return {"dependencies": dependencies}