jaclang 0.6.0__py3-none-any.whl → 0.6.5__py3-none-any.whl

jaclang/langserve/tests/pylsp_jsonrpc/streams.py

@@ -0,0 +1,113 @@
+# Copyright 2017-2020 Palantir Technologies, Inc.
+# Copyright 2021- Python Language Server Contributors.
+
+import logging
+import threading
+
+try:
+    import ujson as json
+except Exception:  # pylint: disable=broad-except
+    import json
+
+log = logging.getLogger(__name__)
+
+
+class JsonRpcStreamReader:
+    def __init__(self, rfile):
+        self._rfile = rfile
+
+    def close(self):
+        self._rfile.close()
+
+    def listen(self, message_consumer):
+        """Blocking call to listen for messages on the rfile.
+
+        Args:
+            message_consumer (fn): function that is passed each message as it is read off the socket.
+        """
+        while not self._rfile.closed:
+            try:
+                request_str = self._read_message()
+            except ValueError:
+                if self._rfile.closed:
+                    return
+                log.exception("Failed to read from rfile")
+
+            if request_str is None:
+                break
+
+            try:
+                message_consumer(json.loads(request_str.decode("utf-8")))
+            except ValueError:
+                log.exception("Failed to parse JSON message %s", request_str)
+                continue
+
+    def _read_message(self):
+        """Reads the contents of a message.
+
+        Returns:
+            body of message if parsable else None
+        """
+        line = self._rfile.readline()
+
+        if not line:
+            return None
+
+        content_length = self._content_length(line)
+
+        # Blindly consume all header lines
+        while line and line.strip():
+            line = self._rfile.readline()
+
+        if not line:
+            return None
+
+        # Grab the body
+        return self._rfile.read(content_length)
+
+    @staticmethod
+    def _content_length(line):
+        """Extract the content length from an input line."""
+        if line.startswith(b"Content-Length: "):
+            _, value = line.split(b"Content-Length: ")
+            value = value.strip()
+            try:
+                return int(value)
+            except ValueError as e:
+                raise ValueError(f"Invalid Content-Length header: {value}") from e
+
+        return None
+
+
+class JsonRpcStreamWriter:
+    def __init__(self, wfile, **json_dumps_args):
+        self._wfile = wfile
+        self._wfile_lock = threading.Lock()
+        self._json_dumps_args = json_dumps_args
+
+    def close(self):
+        with self._wfile_lock:
+            self._wfile.close()
+
+    def write(self, message):
+        with self._wfile_lock:
+            if self._wfile.closed:
+                return
+            try:
+                body = json.dumps(message, **self._json_dumps_args)
+
+                # Ensure we get the byte length, not the character length
+                content_length = (
+                    len(body) if isinstance(body, bytes) else len(body.encode("utf-8"))
+                )
+
+                response = (
+                    f"Content-Length: {content_length}\r\n"
+                    f"Content-Type: application/vscode-jsonrpc; charset=utf8\r\n\r\n"
+                    f"{body}"
+                )
+
+                self._wfile.write(response.encode("utf-8"))
+                self._wfile.flush()
+            except Exception:  # pylint: disable=broad-except
+                log.exception("Failed to write message to output file %s", message)

jaclang/langserve/tests/session.py

@@ -0,0 +1,255 @@
+"""Provides LSP session helpers for testing."""
+
+import os
+import subprocess
+import sys
+from concurrent.futures import Future, ThreadPoolExecutor
+from threading import Event
+
+from .pylsp_jsonrpc.dispatchers import MethodDispatcher
+from .pylsp_jsonrpc.endpoint import Endpoint
+from .pylsp_jsonrpc.streams import JsonRpcStreamReader, JsonRpcStreamWriter
+
+from . import defaults
+
+LSP_EXIT_TIMEOUT = 5000
+
+
+PUBLISH_DIAGNOSTICS = "textDocument/publishDiagnostics"
+WINDOW_LOG_MESSAGE = "window/logMessage"
+WINDOW_SHOW_MESSAGE = "window/showMessage"
+
+
+class LspSession(MethodDispatcher):
+    """Send and Receive messages over LSP as a test LS Client."""
+
+    def __init__(self, cwd=None):
+        self.cwd = cwd if cwd else os.getcwd()
+        self._thread_pool = ThreadPoolExecutor()
+        self._sub = None
+        self._writer = None
+        self._reader = None
+        self._endpoint = None
+        self._notification_callbacks = {}
+
+    def __enter__(self):
+        """Context manager entrypoint.
+
+        shell=True needed for pytest-cov to work in subprocess.
+        """
+        self._sub = subprocess.Popen(
+            [
+                sys.executable,
+                os.path.join(os.path.dirname(os.path.dirname(__file__)), "server.py"),
+            ],
+            stdout=subprocess.PIPE,
+            stdin=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            bufsize=0,
+            cwd=self.cwd,
+            env=os.environ,
+            shell="WITH_COVERAGE" in os.environ,
+        )
+
+        self._writer = JsonRpcStreamWriter(os.fdopen(self._sub.stdin.fileno(), "wb"))
+        self._reader = JsonRpcStreamReader(os.fdopen(self._sub.stdout.fileno(), "rb"))
+
+        dispatcher = {
+            PUBLISH_DIAGNOSTICS: self._publish_diagnostics,
+            WINDOW_SHOW_MESSAGE: self._window_show_message,
+            WINDOW_LOG_MESSAGE: self._window_log_message,
+        }
+        self._endpoint = Endpoint(dispatcher, self._writer.write)
+        self._thread_pool.submit(self._reader.listen, self._endpoint.consume)
+        return self
+
+    def __exit__(self, typ, value, _tb):
+        self.shutdown(True)
+        try:
+            self._sub.terminate()
+        except Exception as e:
+            pass
+        self._endpoint.shutdown()
+        self._thread_pool.shutdown()
+
+    def initialize(
+        self,
+        initialize_params=None,
+        process_server_capabilities=None,
+    ):
+        """Sends the initialize request to LSP server."""
+        server_initialized = Event()
+
+        def _after_initialize(fut):
+            if process_server_capabilities:
+                process_server_capabilities(fut.result())
+            self.initialized()
+            server_initialized.set()
+
+        self._send_request(
+            "initialize",
+            params=(
+                initialize_params
+                if initialize_params is not None
+                else defaults.VSCODE_DEFAULT_INITIALIZE
+            ),
+            handle_response=_after_initialize,
+        )
+
+        server_initialized.wait()
+
+    def initialized(self, initialized_params=None):
+        """Sends the initialized notification to LSP server."""
+        if initialized_params is None:
+            initialized_params = {}
+        self._endpoint.notify("initialized", initialized_params)
+
+    def shutdown(self, should_exit, exit_timeout=LSP_EXIT_TIMEOUT):
+        """Sends the shutdown request to LSP server."""
+
+        def _after_shutdown(_):
+            if should_exit:
+                self.exit_lsp(exit_timeout)
+
+        self._send_request("shutdown", handle_response=_after_shutdown)
+
+    def exit_lsp(self, exit_timeout=LSP_EXIT_TIMEOUT):
+        """Handles LSP server process exit."""
+        self._endpoint.notify("exit")
+        assert self._sub.wait(exit_timeout) == 0
+
+    def text_document_completion(self, completion_params):
+        """Sends text document completion request to LSP server."""
+        fut = self._send_request("textDocument/completion", params=completion_params)
+        return fut.result()
+
+    def text_document_rename(self, rename_params):
+        """Sends text document rename request to LSP server."""
+        fut = self._send_request("textDocument/rename", params=rename_params)
+        return fut.result()
+
+    def text_document_code_action(self, code_action_params):
+        """Sends text document code action request to LSP server."""
+        fut = self._send_request("textDocument/codeAction", params=code_action_params)
+        return fut.result()
+
+    def text_document_hover(self, hover_params):
+        """Sends text document hover request to LSP server."""
+        fut = self._send_request("textDocument/hover", params=hover_params)
+        return fut.result()
+
+    def text_document_formatting(self, formatting_params):
+        """Sends text document formatting request to LSP server."""
+        fut = self._send_request("textDocument/formatting", params=formatting_params)
+        return fut.result()
+
+    def text_document_signature_help(self, signature_help_params):
+        """Sends text document hover request to LSP server."""
+        fut = self._send_request(
+            "textDocument/signatureHelp", params=signature_help_params
+        )
+        return fut.result()
+
+    def text_document_declaration(self, declaration_params):
+        """Sends text document declaration request to LSP server."""
+        fut = self._send_request("textDocument/declaration", params=declaration_params)
+        return fut.result()
+
+    def text_document_definition(self, definition_params):
+        """Sends text document definition request to LSP server."""
+        fut = self._send_request("textDocument/definition", params=definition_params)
+        return fut.result()
+
+    def text_document_symbol(self, document_symbol_params):
+        """Sends text document symbol request to LSP server."""
+        fut = self._send_request(
+            "textDocument/documentSymbol", params=document_symbol_params
+        )
+        return fut.result()
+
+    def text_document_highlight(self, document_highlight_params):
+        """Sends text document highlight request to LSP server."""
+        fut = self._send_request(
+            "textDocument/documentHighlight", params=document_highlight_params
+        )
+        return fut.result()
+
+    def text_document_references(self, references_params):
+        """Sends text document references request to LSP server."""
+        fut = self._send_request("textDocument/references", params=references_params)
+        return fut.result()
+
+    def workspace_symbol(self, workspace_symbol_params):
+        """Sends workspace symbol request to LSP server."""
+        fut = self._send_request("workspace/symbol", params=workspace_symbol_params)
+        return fut.result()
+
+    def completion_item_resolve(self, resolve_params):
+        """Sends completion item resolve request to LSP server."""
+        fut = self._send_request("completionItem/resolve", params=resolve_params)
+        return fut.result()
+
+    def notify_did_change(self, did_change_params):
+        """Sends did change notification to LSP Server."""
+        self._send_notification("textDocument/didChange", params=did_change_params)
+
+    def notify_did_save(self, did_save_params):
+        """Sends did save notification to LSP Server."""
+        self._send_notification("textDocument/didSave", params=did_save_params)
+
+    def notify_did_open(self, did_open_params):
+        """Sends did open notification to LSP Server."""
+        self._send_notification("textDocument/didOpen", params=did_open_params)
+
+    def set_notification_callback(self, notification_name, callback):
+        """Set custom LS notification handler."""
+        self._notification_callbacks[notification_name] = callback
+
+    def get_notification_callback(self, notification_name):
+        """Gets callback if set or default callback for a given LS notification."""
+        try:
+            return self._notification_callbacks[notification_name]
+        except KeyError:
+
+            def _default_handler(_params):
+                """Default notification handler."""
+
+            return _default_handler
+
+    def _publish_diagnostics(self, publish_diagnostics_params):
+        """Internal handler for text document publish diagnostics."""
+        return self._handle_notification(
+            PUBLISH_DIAGNOSTICS, publish_diagnostics_params
+        )
+
+    def _window_log_message(self, window_log_message_params):
+        """Internal handler for window log message."""
+        return self._handle_notification(WINDOW_LOG_MESSAGE, window_log_message_params)
+
+    def _window_show_message(self, window_show_message_params):
+        """Internal handler for window show message."""
+        return self._handle_notification(
+            WINDOW_SHOW_MESSAGE, window_show_message_params
+        )
+
+    def _handle_notification(self, notification_name, params):
+        """Internal handler for notifications."""
+        fut = Future()
+
+        def _handler():
+            callback = self.get_notification_callback(notification_name)
+            callback(params)
+            fut.set_result(None)
+
+        self._thread_pool.submit(_handler)
+        return fut
+
+    def _send_request(self, name, params=None, handle_response=lambda f: f.done()):
+        """Sends {name} request to the LSP server."""
+        fut = self._endpoint.request(name, params)
+        fut.add_done_callback(handle_response)
+        return fut
+
+    def _send_notification(self, name, params=None):
+        """Sends {name} notification to the LSP server."""
+        self._endpoint.notify(name, params)

jaclang/langserve/tests/test_server.py

@@ -0,0 +1,36 @@
+from jaclang.utils.test import TestCase
+from jaclang.vendor.pygls import uris
+from .session import LspSession
+
+
+class TestJacLangServer(TestCase):
+
+    def test_formatting(self) -> None:
+        with LspSession() as s:
+            s.initialize()
+            with open(self.fixture_abs_path("hello.jac"), "w") as f:
+                f.write('with entry {print("Hello, World!");}')
+            with open(self.fixture_abs_path("hello.jac"), "r") as f:
+                self.assertEqual(f.read(), 'with entry {print("Hello, World!");}')
+            params_json = {
+                "textDocument": {
+                    "uri": uris.from_fs_path(self.fixture_abs_path("hello.jac")),
+                    "languageId": "jac",
+                    "version": 1,
+                    "text": "with entry {print('Hello, World!');}",
+                },
+                "options": {"tabSize": 4, "insertSpaces": True},
+            }
+            text_edits = s.text_document_formatting(params_json)
+            self.assertEqual(
+                text_edits,
+                [
+                    {
+                        "range": {
+                            "start": {"line": 0, "character": 0},
+                            "end": {"line": 4, "character": 0},
+                        },
+                        "newText": 'with entry {\n    print("Hello, World!");\n}\n',
+                    }
+                ],
+            )

jaclang/langserve/utils.py

@@ -0,0 +1,53 @@
+"""Utility functions for the language server."""
+
+import asyncio
+import logging
+from functools import wraps
+from typing import Any, Awaitable, Callable, Coroutine, ParamSpec, TypeVar
+
+from jaclang.vendor.pygls.server import LanguageServer
+
+import lsprotocol.types as lspt
+
+
+def log_error(ls: LanguageServer, message: str) -> None:
+    """Log an error message."""
+    ls.show_message_log(message, lspt.MessageType.Error)
+    ls.show_message(message, lspt.MessageType.Error)
+
+
+def log(info: str) -> None:
+    """Log an info message."""
+    logging.warning(info)
+
+
+T = TypeVar("T", bound=Callable[..., Coroutine[Any, Any, Any]])
+P = ParamSpec("P")
+
+
+def debounce(wait: float) -> Callable[[T], Callable[..., Awaitable[None]]]:
+    """Debounce decorator for async functions."""
+
+    def decorator(fn: T) -> Callable[..., Awaitable[None]]:
+        @wraps(fn)
+        async def debounced(*args: P.args, **kwargs: P.kwargs) -> None:
+            async def call_it() -> None:
+                await fn(*args, **kwargs)
+
+            if hasattr(debounced, "_task"):
+                debounced._task.cancel()
+
+            async def debounced_coro() -> None:
+                try:
+                    await asyncio.sleep(wait)
+                    await call_it()
+                except asyncio.CancelledError:
+                    pass
+
+            setattr(  # noqa: B010
+                debounced, "_task", asyncio.create_task(debounced_coro())
+            )
+
+        return debounced
+
+    return decorator

jaclang/plugin/__init__.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from .default import hookimpl
-from .spec import Architype, DSFunc
+from .default import Architype, DSFunc, hookimpl
 
 __all__ = ["Architype", "DSFunc", "hookimpl"]

jaclang/plugin/builtin.py

@@ -20,9 +20,9 @@ def dotgen(
     dot_file: Optional[str] = None,
 ) -> str:
     """Print the dot graph."""
-    from jaclang.core.construct import root
     from jaclang.plugin.feature import pm
 
+    root = pm.hook.get_root()
     node = node if node is not None else root
     depth = depth if depth is not None else -1
     traverse = traverse if traverse is not None else False

jaclang/plugin/default.py

@@ -7,9 +7,11 @@ import html
 import os
 import pickle
 import types
+from contextvars import ContextVar
 from dataclasses import field
 from functools import wraps
 from typing import Any, Callable, Optional, Type, Union
+from uuid import UUID
 
 from jaclang.compiler.absyntree import Module
 from jaclang.compiler.constant import EdgeDir, colors
@@ -34,10 +36,11 @@ from jaclang.core.construct import (
     Root,
     WalkerAnchor,
     WalkerArchitype,
-    root,
 )
 from jaclang.core.importer import jac_importer
+from jaclang.core.memory import Memory
 from jaclang.core.registry import SemInfo, SemRegistry, SemScope
+from jaclang.core.shelve_storage import ShelveStorage
 from jaclang.core.utils import traverse_graph
 from jaclang.plugin.feature import JacFeature as Jac
 from jaclang.plugin.spec import T
@@ -55,11 +58,10 @@ __all__ = [
     "WalkerAnchor",
     "NodeArchitype",
     "EdgeArchitype",
+    "Root",
     "WalkerArchitype",
     "Architype",
     "DSFunc",
-    "root",
-    "Root",
     "jac_importer",
     "T",
 ]
@@ -68,11 +70,95 @@ __all__ = [
 hookimpl = pluggy.HookimplMarker("jac")
 
 
+class ExecutionContext:
+    """Default Execution Context implementation."""
+
+    mem: Optional[Memory]
+    root: Optional[Root]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.mem = ShelveStorage()
+        self.root = None
+
+    def init_memory(self, session: str = "") -> None:
+        if session:
+            self.mem = ShelveStorage(session)
+        else:
+            self.mem = Memory()
+
+    def get_root(self) -> Root:
+        if self.mem is None:
+            raise ValueError("Memory not initialized")
+
+        if not self.root:
+            root = self.mem.get_obj(UUID(int=0))
+            if root is None:
+                self.root = Root()
+                self.mem.save_obj(self.root, persistent=self.root._jac_.persistent)
+            elif not isinstance(root, Root):
+                raise ValueError(f"Invalid root object: {root}")
+            else:
+                self.root = root
+        return self.root
+
+    def get_obj(self, obj_id: UUID) -> Architype | None:
+        """Get object from memory."""
+        if self.mem is None:
+            raise ValueError("Memory not initialized")
+
+        return self.mem.get_obj(obj_id)
+
+    def save_obj(self, item: Architype, persistent: bool) -> None:
+        """Save object to memory."""
+        if self.mem is None:
+            raise ValueError("Memory not initialized")
+
+        self.mem.save_obj(item, persistent)
+
+    def reset(self) -> None:
+        """Reset the execution context."""
+        if self.mem:
+            self.mem.close()
+        self.mem = None
+        self.root = None
+
+
+ExecContext: ContextVar[ExecutionContext | None] = ContextVar(
+    "ExecutionContext", default=None
+)
+
+
 class JacFeatureDefaults:
     """Jac Feature."""
 
     pm = pluggy.PluginManager("jac")
 
+    @staticmethod
+    @hookimpl
+    def context(session: str = "") -> ExecutionContext:
+        """Get the execution context."""
+        ctx = ExecContext.get()
+        if ctx is None:
+            ctx = ExecutionContext()
+            ExecContext.set(ctx)
+        return ctx
+
+    @staticmethod
+    @hookimpl
+    def reset_context() -> None:
+        """Reset the execution context."""
+        ctx = ExecContext.get()
+        if ctx:
+            ctx.reset()
+        ExecContext.set(None)
+
+    @staticmethod
+    @hookimpl
+    def memory_hook() -> Memory | None:
+        """Return the memory hook."""
+        return Jac.context().mem
+
     @staticmethod
     @hookimpl
     def make_architype(
@@ -377,6 +463,12 @@ class JacFeatureDefaults:
                 conn_edge = edge_spec()
                 edges.append(conn_edge)
                 i._jac_.connect_node(j, conn_edge)
+
+                if i._jac_.persistent or j._jac_.persistent:
+                    conn_edge.save()
+                    j.save()
+                    i.save()
+
         return right if not edges_only else edges
 
     @staticmethod
@@ -429,7 +521,7 @@ class JacFeatureDefaults:
     @hookimpl
     def get_root() -> Root:
         """Jac's assign comprehension feature."""
-        return root
+        return Jac.context().get_root()
 
     @staticmethod
     @hookimpl
@@ -603,6 +695,7 @@ class JacFeatureDefaults:
         )
 
         type_collector: list = []
+        incl_info = [x for x in incl_info if not isinstance(x[1], type)]
         information, collected_types = get_info_types(_scope, mod_registry, incl_info)
         type_collector.extend(collected_types)
         inputs_information_list = []

jaclang/plugin/feature.py

@@ -3,18 +3,22 @@
 from __future__ import annotations
 
 import types
-from typing import Any, Callable, Optional, Type, Union
+from typing import Any, Callable, Optional, TYPE_CHECKING, Type, Union
 
 from jaclang.compiler.absyntree import Module
-from jaclang.core.construct import (
-    Architype,
-    EdgeArchitype,
-    NodeArchitype,
-    Root,
-    WalkerArchitype,
-)
 from jaclang.plugin.spec import JacBuiltin, JacCmdSpec, JacFeatureSpec, T
 
+if TYPE_CHECKING:
+    from jaclang.core.construct import (
+        Architype,
+        EdgeArchitype,
+        NodeArchitype,
+        WalkerArchitype,
+        Root,
+    )
+    from jaclang.plugin.default import ExecutionContext
+    from jaclang.core.memory import Memory
+
 
 import pluggy
 
@@ -28,8 +32,23 @@ class JacFeature:
     """Jac Feature."""
 
     import abc
-    from jaclang.plugin.spec import DSFunc
     from jaclang.compiler.constant import EdgeDir
+    from jaclang.plugin.spec import DSFunc
+
+    @staticmethod
+    def context(session: str = "") -> ExecutionContext:
+        """Create execution context."""
+        return pm.hook.context(session=session)
+
+    @staticmethod
+    def reset_context() -> None:
+        """Reset execution context."""
+        return pm.hook.reset_context()
+
+    @staticmethod
+    def memory_hook() -> Memory | None:
+        """Create memory abstraction."""
+        return pm.hook.memory_hook()
 
     @staticmethod
     def make_architype(