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

setta/lsp/reader_fns/documentHighlight.py

@@ -0,0 +1,25 @@
+from setta.code_gen.python.position_line_col import line_col_to_position_batch
+
+
+def process_document_highlight_response(response, metadata):
+    content = []
+    # if it's not None, then it's a variable or import or something like that
+    if response["result"]:
+        # the response doesn't include uri for some reason,
+        # so we keep track of it manually
+        code = metadata["code"]
+
+        for x in response["result"]:
+            r = x["range"]
+            positions = line_col_to_position_batch(
+                code,
+                [
+                    (r["start"]["line"], r["start"]["character"]),
+                    (r["end"]["line"], r["end"]["character"]),
+                ],
+            )
+            content.append({"startPos": positions[0], "endPos": positions[1]})
+    return_id = metadata["return_id"]
+    num_messages = metadata["num_messages"]
+    content = {"positions": content, "numMessages": num_messages}
+    return content, return_id

setta/lsp/reader_fns/references.py

@@ -0,0 +1,34 @@
+from setta.code_gen.python.position_line_col import line_col_to_position_batch
+
+
+def process_references_response(response, metadata):
+    content = []
+    if response["result"] and len(response["result"]) > 0:
+        code = metadata["code"]
+        line_cols = []
+        for x in response["result"]:
+            line_cols.append(
+                (
+                    x["range"]["start"]["line"],
+                    x["range"]["start"]["character"],
+                )
+            )
+            line_cols.append(
+                (x["range"]["end"]["line"], x["range"]["end"]["character"])
+            )
+        positions = line_col_to_position_batch(code, line_cols)
+        position_offset = metadata["position_offset"]
+        for idx in range(0, len(positions), 2):
+            startPos = positions[idx] - position_offset
+            endPos = positions[idx + 1] - position_offset
+            content.append(
+                {
+                    "startPos": startPos,
+                    "endPos": endPos,
+                }
+            )
+
+    return_id = metadata["return_id"]
+    num_messages = metadata["num_messages"]
+    content = {"positions": content, "numMessages": num_messages}
+    return content, return_id

setta/lsp/reader_fns/signatureHelp.py

@@ -0,0 +1,99 @@
+import re
+
+import docstring_parser
+
+
+def process_signature_help_response(response):
+    content = None
+    if response["result"]:
+        content = parse_signature(response["result"]["signatures"][0])
+    return content
+
+
+def parse_signature(signature):
+    docstring = process_docstring(signature)
+    param_info_list = wrapper_parse_signature(
+        signature, docstring, get_pyright_param_proposals
+    )
+    if len(param_info_list) == 0:
+        param_info_list = wrapper_parse_signature(
+            signature, docstring, get_docstring_parser_param_proposals
+        )
+    documentation = signature.get("documentation", {"value": ""})["value"]
+    return [documentation, param_info_list]
+
+
+def process_docstring(signature):
+    if "documentation" in signature:
+        parsed = docstring_parser.parse(signature["documentation"]["value"])
+        return {
+            extract_variable_name(v.arg_name): {
+                "default": v.default,
+                "description": v.description,
+            }
+            for v in parsed.params
+        }
+    return {}
+
+
+def wrapper_parse_signature(signature, docstring, get_proposed_params):
+    params = get_proposed_params(signature, docstring)
+    param_info_list = []
+    positional_only_indicator_idx = -1
+
+    for idx, (name, default, description) in enumerate(params):
+        if name == "/":
+            positional_only_indicator_idx = idx
+            continue
+        if name in [None, "", "*", "..."]:
+            continue
+
+        if not default:
+            default = ""
+
+        if not description:
+            description = docstring.get(name, {}).get("description", "")
+
+        # Append parsed information to list
+        param_info_list.append(
+            {
+                "name": name,
+                "defaultVal": default,
+                "description": description,
+                "positionalOnly": False,
+            }
+        )
+
+    for idx in range(0, positional_only_indicator_idx):
+        param_info_list[idx]["positionalOnly"] = True
+
+    return param_info_list
+
+
+def get_docstring_parser_param_proposals(_, docstring):
+    return [(k, v["default"], v["description"]) for k, v in docstring.items()]
+
+
+def get_pyright_param_proposals(signature, _):
+    parameters = signature["parameters"]
+    params = [(p["label"], p["documentation"]["value"]) for p in parameters]
+    proposals = []
+    for param, description in params:
+        # Extract parameter name and default value
+        if param.startswith("/"):
+            proposals.append((param, None, None))
+            continue
+        parts = param.split("=")
+        name = extract_variable_name(parts[0].strip())
+        default = parts[1].strip() if len(parts) > 1 else None
+        proposals.append((name, default, description))
+
+    return proposals
+
+
+def extract_variable_name(code_string):
+    # Find potential variable names
+    potential_names = re.findall(r"\b[a-zA-Z_]\w*\b", code_string)
+
+    # Return the first valid identifier
+    return next((name for name in potential_names if name.isidentifier()), None)

setta/lsp/server.py

@@ -0,0 +1,150 @@
+import asyncio
+import logging
+from collections import defaultdict
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+class LanguageServer:
+    def __init__(self, workspace_folder, code_folder, settings, name):
+        # Create folders if they don't exist
+        workspace_folder.mkdir(parents=True, exist_ok=True)
+        code_folder.mkdir(parents=True, exist_ok=True)
+        self.workspace_folder = workspace_folder
+        self.code_folder = code_folder
+        self.init_properties()
+        self.server = None
+        if not isinstance(settings, dict):
+            logger.warning(
+                "warning: your language server settings is not an object, so it will have no effect"
+            )
+            self.settings = {}
+        else:
+            self.settings = settings
+        self.name = name
+
+    async def start_server(self):
+        logger.debug(f"Starting {self.name} language server")
+        self.init_properties()
+        self.server = await asyncio.create_subprocess_shell(
+            "basedpyright-langserver --stdio",
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+
+    def init_properties(self):
+        self.cleanup_pending_requests()
+        self.pending_requests = {}
+        self.msg_id_to_metadata = defaultdict(dict)
+        self.code_metadata = {}
+        self.initialization_request_id = None
+        self.is_initialized = False
+        self.temp_files = {}
+
+    def cleanup_pending_requests(self):
+        if not hasattr(self, "pending_requests"):
+            return
+        # Cancel all pending futures with an appropriate error
+        for future in self.pending_requests.values():
+            if not future.done():
+                future.set_exception(
+                    ConnectionResetError("Language server was restarted")
+                )
+        self.pending_requests.clear()
+
+    def kill_server(self):
+        try:
+            logger.debug(f"Stopping {self.name} language server")
+            self.server.terminate()
+        except ProcessLookupError:
+            pass
+        for path in self.temp_files.values():
+            if path.exists():
+                path.unlink()
+
+    @property
+    def stdin(self):
+        return self.server.stdin
+
+    @property
+    def stdout(self):
+        return self.server.stdout
+
+
+class LanguageServerInteractor:
+    @property
+    def settings(self):
+        return self.server.settings
+
+    @property
+    def msg_id_to_metadata(self):
+        return self.server.msg_id_to_metadata
+
+    @property
+    def workspace_folder(self):
+        return self.server.workspace_folder
+
+    @property
+    def code_folder(self):
+        return self.server.code_folder
+
+    @property
+    def pending_requests(self):
+        return self.server.pending_requests
+
+    @property
+    def temp_files(self):
+        return self.server.temp_files
+
+    async def wait_for_response(self, id, timeout=None):
+        future = self.pending_requests.get(id)
+        if future:
+            try:
+                result = await asyncio.wait_for(future, timeout=timeout)
+                return result
+            except ConnectionResetError:
+                # Handle the error case
+                pass
+            finally:
+                self.pending_requests.pop(id, None)
+        return None
+
+    def project_code_folder(self, projectConfigName):
+        return self.code_folder / projectConfigName
+
+    def construct_code_id(self, projectConfigName, sectionName):
+        return str(Path(projectConfigName) / sectionName)
+
+    @property
+    def code_metadata(self):
+        return self.server.code_metadata
+
+    def get_new_version_number(
+        self,
+        uri,
+        code,
+        codeJson,
+        ref_template_var_positions,
+        needsDiagnostics,
+        projectConfigId,
+        codeSectionId,
+    ):
+        obj = {
+            "code": code,
+            "codeJson": codeJson,
+            "ref_template_var_positions": ref_template_var_positions,
+            "needsDiagnostics": needsDiagnostics,
+            "projectConfigId": projectConfigId,
+            "codeSectionId": codeSectionId,
+        }
+        if uri not in self.code_metadata:
+            obj["version"] = 0
+        else:
+            obj["version"] = self.code_metadata[uri]["version"] + 1
+        self.code_metadata[uri] = obj
+        return obj["version"]
+
+    def get_code_metadata(self, uri):
+        return self.code_metadata[uri.lower()]

setta/lsp/utils.py

@@ -0,0 +1,60 @@
+from .file_watcher import LSPFileWatcher
+from .reader import LanguageServerReader
+from .server import LanguageServer
+from .writer import LanguageServerWriter
+
+
+def create_lsps(
+    workspace_folder,
+    code_folder,
+    settings,
+):
+    return {
+        "full": LanguageServer(
+            workspace_folder=workspace_folder,
+            code_folder=code_folder,
+            settings=settings,
+            name="full",
+        ),
+        "basic": LanguageServer(
+            workspace_folder=code_folder,
+            code_folder=code_folder,
+            settings={},
+            name="basic",
+        ),
+    }
+
+
+def create_lsp_writers(lsps):
+    return {
+        k: LanguageServerWriter(v, use_virtual_files=k == "basic")
+        for k, v in lsps.items()
+    }
+
+
+def create_lsp_readers(lsps, websocket_manager):
+    return {k: LanguageServerReader(v, websocket_manager) for k, v in lsps.items()}
+
+
+def create_file_watcher(lsps, lsp_writers):
+    file_watcher = LSPFileWatcher(lsp_writers["full"].send_did_change_watched_files)
+    file_watcher.add_workspace_folder(lsps["full"].workspace_folder)
+    return file_watcher
+
+
+async def start_lsps(lsps, lsp_readers, lsp_writers):
+    for k, v in lsps.items():
+        await v.start_server()
+        await lsp_readers[k].start_listener()
+        await lsp_writers[k].send_initialize_request()
+
+
+async def kill_lsps(lsps, lsp_readers):
+    for k, v in lsps.items():
+        await lsp_readers[k].stop()  # Wait for reader to finish
+        v.kill_server()
+
+
+async def restart_lsps(lsps, lsp_readers, lsp_writers):
+    await kill_lsps(lsps, lsp_readers)
+    await start_lsps(lsps, lsp_readers, lsp_writers)

setta/lsp/writer.py

@@ -0,0 +1,306 @@
+import asyncio
+import json
+import logging
+import platform
+import uuid
+
+from setta.utils.utils import recursive_dict_merge
+
+from .server import LanguageServerInteractor
+
+logger = logging.getLogger(__name__)
+
+
+class LanguageServerWriter(LanguageServerInteractor):
+    def __init__(self, server, use_virtual_files):
+        self.server = server
+        self.use_virtual_files = use_virtual_files
+
+    async def send_request(self, method, params, id=None):
+        request_dict = {"jsonrpc": "2.0", "method": method, "params": params}
+        if id:
+            request_dict["id"] = id
+            self.msg_id_to_metadata[id]["method"] = method
+
+        request = json.dumps(request_dict).encode("utf-8")
+        content_length = len(request)
+        header = f"Content-Length: {content_length}\r\n\r\n".encode("utf-8")
+        self.server.stdin.write(header + request)
+        await self.server.stdin.drain()
+        logger.debug(f"Sent {method} with params: {params}")
+
+    async def send_initialize_request(self):
+        params = {
+            "capabilities": {
+                "workspace": {"didChangeWatchedFiles": {"dynamicRegistration": True}}
+            },
+            "workspaceFolders": [{"uri": self.get_workspace_uri()}],
+        }
+        request_id = str(uuid.uuid4())
+        self.server.initialization_request_id = request_id
+        await self.send_request("initialize", params=params, id=request_id)
+        await self.send_didChangeConfiguration()  # pyright requires this step
+
+    async def send_document_open_request(
+        self,
+        code_id,
+        code,
+        codeJson=None,
+        ref_template_var_positions=None,
+        needsDiagnostics=False,
+        projectConfigId=None,
+        codeSectionId=None,
+        virtual=None,
+    ):
+        virtual = self.get_is_virtual(virtual)
+        uri = self.get_uri(code_id, virtual)
+        if uri not in self.code_metadata:
+            if not virtual:
+                self.create_or_update_temp_file(code_id, code)
+            params = {
+                "textDocument": {
+                    "uri": uri,
+                    "languageId": "python",
+                    "version": self.get_new_version_number(
+                        uri,
+                        code,
+                        codeJson,
+                        ref_template_var_positions,
+                        needsDiagnostics,
+                        projectConfigId,
+                        codeSectionId,
+                    ),
+                    "text": code,
+                }
+            }
+            logger.debug(f"send_document_open_request {uri}")
+            await self.send_request("textDocument/didOpen", params=params)
+        else:
+            logger.debug(f"document already open, updating it {uri}")
+            await self.send_document_update(
+                code_id=code_id,
+                code=code,
+                codeJson=codeJson,
+                ref_template_var_positions=ref_template_var_positions,
+                needsDiagnostics=needsDiagnostics,
+                projectConfigId=projectConfigId,
+                codeSectionId=codeSectionId,
+                virtual=virtual,
+            )
+
+    async def send_autocomplete_request(
+        self,
+        line,
+        character,
+        code_id,
+        request_id,
+        referencable_var_names,
+        generated_var_names,
+        cursor_offset,
+        otherData=None,
+        virtual=None,
+    ):
+        virtual = self.get_is_virtual(virtual)
+        uri = self.get_uri(code_id, virtual)
+        params = {
+            "textDocument": {"uri": uri},
+            "position": {"line": line, "character": character},
+        }
+        self.msg_id_to_metadata[request_id][
+            "referencable_var_names"
+        ] = referencable_var_names
+        self.msg_id_to_metadata[request_id]["generated_var_names"] = generated_var_names
+        self.msg_id_to_metadata[request_id]["cursor_offset"] = cursor_offset
+        self.msg_id_to_metadata[request_id]["otherData"] = otherData
+        self.msg_id_to_metadata[request_id]["code"] = self.get_code_metadata(uri)[
+            "code"
+        ]
+        await self.send_request(
+            "textDocument/completion",
+            params=params,
+            id=request_id,
+        )
+
+    async def send_signature_help_request(
+        self, line, character, code_id, request_id, otherData=None, virtual=None
+    ):
+        virtual = self.get_is_virtual(virtual)
+        uri = self.get_uri(code_id, virtual)
+        params = {
+            "textDocument": {"uri": uri},
+            "position": {"line": line, "character": character},
+            "context": {"triggerKind": 2, "triggerCharacter": "("},
+        }
+        self.msg_id_to_metadata[request_id]["otherData"] = otherData
+        await self.send_request(
+            "textDocument/signatureHelp",
+            params=params,
+            id=request_id,
+        )
+
+    async def send_document_highlight_request(
+        self,
+        line,
+        character,
+        code_id,
+        request_id,
+        return_id,
+        num_messages,
+        virtual=None,
+    ):
+        virtual = self.get_is_virtual(virtual)
+        uri = self.get_uri(code_id, virtual)
+        params = {
+            "textDocument": {"uri": uri},
+            "position": {"line": line, "character": character},
+        }
+        self.msg_id_to_metadata[request_id]["return_id"] = return_id
+        self.msg_id_to_metadata[request_id]["num_messages"] = num_messages
+        self.msg_id_to_metadata[request_id]["code"] = self.get_code_metadata(uri)[
+            "code"
+        ]
+        await self.send_request(
+            "textDocument/documentHighlight",
+            params=params,
+            id=request_id,
+        )
+
+    async def send_didChangeConfiguration(self):
+        # THIS FORMAT WORKS
+        # hardcoded_settings = {
+        #     "basedpyright": {"analysis": {"exclude": ["setta_code"]}}
+        # }
+        hardcoded_settings = {}
+        settings = recursive_dict_merge(hardcoded_settings, self.settings)
+
+        await self.send_request(
+            "workspace/didChangeConfiguration",
+            params={"settings": settings},
+        )
+
+    async def send_document_update(
+        self,
+        code_id,
+        code,
+        codeJson,
+        ref_template_var_positions=None,
+        needsDiagnostics=False,
+        projectConfigId=None,
+        codeSectionId=None,
+        virtual=None,
+    ):
+        virtual = self.get_is_virtual(virtual)
+        if not virtual:
+            self.create_or_update_temp_file(code_id, code)
+        uri = self.get_uri(code_id, virtual)
+        params = {
+            "textDocument": {
+                "uri": uri,
+                "version": self.get_new_version_number(
+                    uri,
+                    code,
+                    codeJson,
+                    ref_template_var_positions,
+                    needsDiagnostics,
+                    projectConfigId,
+                    codeSectionId,
+                ),
+            },
+            "contentChanges": [{"text": code}],
+        }
+        await self.send_request("textDocument/didChange", params=params)
+
+    async def send_reference_request(
+        self,
+        position_offset,
+        line,
+        character,
+        code_id,
+        request_id,
+        return_id,
+        num_messages,
+        virtual=None,
+    ):
+        virtual = self.get_is_virtual(virtual)
+        uri = self.get_uri(code_id, virtual)
+        params = {
+            "textDocument": {
+                "uri": uri,
+            },
+            "position": {"line": line, "character": character},
+            "context": {"includeDeclaration": False},
+        }
+        self.msg_id_to_metadata[request_id]["position_offset"] = position_offset
+        self.msg_id_to_metadata[request_id]["return_id"] = return_id
+        self.msg_id_to_metadata[request_id]["num_messages"] = num_messages
+        self.msg_id_to_metadata[request_id]["code"] = self.get_code_metadata(uri)[
+            "code"
+        ]
+        await self.send_request("textDocument/references", params=params, id=request_id)
+
+    async def send_definition_request(
+        self,
+        line,
+        character,
+        code_id,
+        request_id,
+        virtual=None,
+    ):
+        virtual = self.get_is_virtual(virtual)
+        uri = self.get_uri(code_id, virtual)
+        params = {
+            "textDocument": {"uri": uri},
+            "position": {"line": line, "character": character},
+        }
+        await self.send_request(
+            "textDocument/definition",
+            params=params,
+            id=request_id,
+        )
+
+    async def send_did_change_watched_files(self, changes):
+        params = {"changes": changes}
+        await self.send_request("workspace/didChangeWatchedFiles", params=params)
+
+    def create_asyncio_task(self, async_task, request_id, local=False):
+        if local:
+            self.pending_requests[request_id] = asyncio.Future()
+        asyncio.create_task(async_task)
+
+    def get_workspace_uri(self):
+        return f"file://{self.workspace_folder}".lower()
+
+    def get_base_uri(self, code_id):
+        return self.code_folder / f"{code_id}.py"
+
+    # https://github.com/microsoft/language-server-protocol/issues/1676
+    def get_uri(self, code_id, virtual):
+        base_uri = self.get_base_uri(code_id)
+        prefix = "untitled:" if virtual else "file://"
+        uri = f"{prefix}{base_uri}".lower()
+        if platform.system() == "Windows":
+            uri += "/"  # need trailing slash on Windows for some reason
+        return uri
+
+    def get_is_virtual(self, virtual):
+        return virtual if virtual is not None else self.use_virtual_files
+
+    def create_temp_file(self, code_id, code):
+        """Create a temporary file for the given code and return its path."""
+        # Create all necessary parent directories
+        file_path = self.get_base_uri(code_id)
+        logger.debug(f"Creating temporary file {file_path}")
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write the code to the file
+        file_path.write_text(code)
+
+        # Store the path for later cleanup
+        self.temp_files[code_id] = file_path
+        return file_path
+
+    def create_or_update_temp_file(self, code_id, code):
+        if code_id in self.temp_files:
+            self.temp_files[code_id].write_text(code)
+        else:
+            self.create_temp_file(code_id, code)

setta/routers/__init__.py

@@ -0,0 +1,11 @@
+from .artifact import router as artifact_router
+from .code_info import router as code_info_router
+from .in_memory_fn_stdout_websocket import router as in_memory_fn_stdout_router
+from .interactive import router as interactive_code_router
+from .lsp import router as lsp_router
+from .projects import router as projects_router
+from .reference_renaming import router as reference_renaming_router
+from .sections import router as sections_router
+from .settings import router as settings_router
+from .terminals import router as terminals_router
+from .websocket import router as websocket_router