• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3#
4# Copyright (c) 2024-2025 Huawei Device Co., Ltd.
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18from collections.abc import AsyncIterator
19from contextlib import AbstractAsyncContextManager, asynccontextmanager
20from dataclasses import dataclass
21from signal import SIGTERM
22from typing import Protocol
23
24import trio
25from pytest import fixture
26
27from arkdb import layouts
28from arkdb.compiler import StringCodeCompiler
29from arkdb.debug_client import DebuggerClient, DebugLocator
30from arkdb.debug_types import Locator, Paused, paused_locator
31from arkdb.logs import RichLogger, logger
32from arkdb.runnable_module import ScriptFile
33from arkdb.runtime import DEFAULT_ENTRY_POINT, Runtime, RuntimeProcess
34from arkdb.source_meta import SourceMeta, parse_source_meta
35
36LOG = logger(__name__)
37
38
39class CompileAndRunType(Protocol):
40    def __call__(
41        self,
42        code_or_file: str | ScriptFile,
43        *,
44        entry_point: str = DEFAULT_ENTRY_POINT,
45    ) -> AbstractAsyncContextManager[RuntimeProcess]:
46        pass
47
48
49@fixture
50def compile_and_run(
51    code_compiler: StringCodeCompiler,
52    ark_runtime: Runtime,
53    nursery: trio.Nursery,
54) -> CompileAndRunType:
55    """
56    Return a :class:`CompileAndRunType` function that compiles and executes ``code_or_file`` script.
57    """
58
59    @asynccontextmanager
60    async def run(
61        code_or_file: str | ScriptFile,
62        *,
63        entry_point: str = DEFAULT_ENTRY_POINT,
64    ) -> AsyncIterator[RuntimeProcess]:
65        script_file = code_compiler.compile(source_code=code_or_file) if isinstance(code_or_file, str) else code_or_file
66        meta = parse_source_meta(script_file.read_text())
67        script_file.log(LOG, highlight_lines=[bp.line_number for bp in meta.breakpoints])
68        async with ark_runtime.run(
69            nursery,
70            module=script_file,
71            entry_point=entry_point,
72            debug=False,
73        ) as process:
74            yield process
75            await process.wait()
76
77    return run
78
79
80@dataclass
81class DebugContext:
82    script_file: ScriptFile
83    client: DebuggerClient
84    process: RuntimeProcess
85    meta: SourceMeta
86
87
88class CompileAndResumeType(Protocol):
89    def __call__(
90        self,
91        code_or_file: str | ScriptFile,
92        *,
93        entry_point: str = DEFAULT_ENTRY_POINT,
94    ) -> AbstractAsyncContextManager[DebugContext]:
95        pass
96
97
98@fixture
99def debug_locator(code_compiler: StringCodeCompiler) -> DebugLocator:
100    """
101    Return a :class:`DebugLocator` instance.
102    """
103    return DebugLocator(code_compiler, url="ws://localhost:19015")
104
105
106@asynccontextmanager
107async def _connect_and_set_breakpoints(
108    nursery: trio.Nursery,
109    debug_locator: DebugLocator,
110    meta: SourceMeta,
111    script_file: ScriptFile,
112    process: RuntimeProcess,
113):
114    try:
115        async with debug_locator.connect(nursery) as client:
116            await client.configure(nursery)
117            await client.run_if_waiting_for_debugger()
118            for bp in meta.breakpoints:
119                bp.breakpoint_id, bp.locations = await client.set_breakpoint_by_url(
120                    url=script_file.source_file.name,
121                    line_number=bp.line_number,
122                )
123            yield DebugContext(
124                script_file=script_file,
125                client=client,
126                process=process,
127                meta=meta,
128            )
129            await trio.lowlevel.checkpoint()
130    finally:
131        process.terminate()
132        with trio.CancelScope(deadline=1, shield=True):
133            await process.wait()
134
135
136@fixture
137def compile_and_resume(
138    code_compiler: StringCodeCompiler,
139    ark_runtime: Runtime,
140    nursery: trio.Nursery,
141    debug_locator: DebugLocator,
142) -> CompileAndResumeType:
143
144    @asynccontextmanager
145    async def run(
146        code_or_file: str | ScriptFile,
147        *,
148        entry_point: str = DEFAULT_ENTRY_POINT,
149    ) -> AsyncIterator[DebugContext]:
150        script_file = code_compiler.compile(source_code=code_or_file) if isinstance(code_or_file, str) else code_or_file
151        meta = parse_source_meta(script_file.read_text())
152        script_file.log(LOG, highlight_lines=[bp.line_number for bp in meta.breakpoints])
153        async with ark_runtime.run(
154            nursery,
155            module=script_file,
156            entry_point=entry_point,
157            debug=True,
158        ) as process:
159            async with _connect_and_set_breakpoints(nursery, debug_locator, meta, script_file, process) as context:
160                yield context
161        if process.returncode != -SIGTERM:
162            raise RuntimeError()
163
164    return run
165
166
167class StopOnPausedType(Protocol):
168    def __call__(
169        self,
170        code_or_file: str | ScriptFile,
171        *,
172        entry_point: str = ...,
173    ) -> AbstractAsyncContextManager[Paused]:
174        pass
175
176
177@fixture
178def run_and_stop_on_breakpoint(
179    compile_and_resume: CompileAndResumeType,
180    log: RichLogger,
181) -> StopOnPausedType:
182    """
183    Return a :class:`StopOnPausedType` function that return :class:`debug_types.Paused` context.
184    """
185
186    @asynccontextmanager
187    async def run(
188        code_or_file: str | ScriptFile,
189        *,
190        entry_point: str = DEFAULT_ENTRY_POINT,
191    ) -> AsyncIterator[Paused]:
192        async with compile_and_resume(code_or_file, entry_point=entry_point) as context:
193            paused = await context.client.resume_and_wait_for_paused()
194            log.info(
195                context.script_file.source_file,
196                rich=await layouts.paused_layout(
197                    paused=paused_locator(paused=paused, client=context.client, meta=context.meta),
198                    url=context.script_file.source_file,
199                ),
200            )
201            yield Paused(locator=Locator(client=context.client), data=paused)
202
203    return run
204