1'''
2This module contains functions to add custom scripts, which can be embedded
3into logic programs.
4
5Examples
6--------
7The following example shows how to add a script that works the same way as
8clingo's embedded Python script:
9```python
10>>> from clingo.script import Script, register_script
11>>> from clingo.control import Control
12>>>
13>>> import __main__
14>>>
15>>> class MyPythonScript(Script):
16...     def execute(self, location, code):
17...         exec(code, __main__.__dict__, __main__.__dict__)
18...     def call(self, location, name, arguments):
19...         return getattr(__main__, name)(*arguments)
20...     def callable(self, name):
21...         return name in __main__.__dict__ and callable(__main__.__dict__[name])
22...
23>>> register_script('mypython', MyPythonScript())
24>>>
25>>> ctl = Control()
26>>> ctl.add('base', [], """
27... #script(mypython)
28... from clingo.symbol import Number
29... def func(a):
30...     return Number(a.number + 1)
31... #end.
32... a(@func(1)).
33... """)
34>>>
35>>> ctl.ground([('base',[])])
36>>> ctl.solve(on_model=print)
37a(2)
38```
39'''
40
41from platform import python_version
42from abc import ABCMeta, abstractmethod
43from typing import Any, List, Iterable, Tuple, Union
44from collections.abc import Iterable as IterableABC
45from traceback import format_exception
46from clingo._internal import _c_call, _ffi, _handle_error, _lib
47from clingo.control import Control
48from clingo.symbol import Symbol
49from clingo.ast import Location, _py_location
50
51try:
52    import __main__ # type: ignore
53except ImportError:
54    # Note: pypy does not create a main module if embedded
55    import sys
56    import types
57    sys.modules['__main__'] = types.ModuleType('__main__', 'the main module')
58    import __main__ # type: ignore
59
60__all__ = [ 'Script', 'enable_python', 'register_script' ]
61
62def _cb_error_top_level(exception, exc_value, traceback):
63    msg = "".join(format_exception(exception, exc_value, traceback))
64    _lib.clingo_set_error(_lib.clingo_error_runtime, msg.encode())
65    return False
66
67
68class Script(metaclass=ABCMeta):
69    '''
70    This interface can be implemented to embed custom scripting languages into
71    logic programs.
72    '''
73    @abstractmethod
74    def execute(self, location: Location, code: str) -> None:
75        '''
76        Execute the given source code.
77
78        Parameters
79        ----------
80        location
81            The location of the code.
82        code
83            The code to execute.
84        '''
85
86    @abstractmethod
87    def call(self, location: Location, name: str, arguments: Iterable[Symbol]) -> Union[Iterable[Symbol], Symbol]:
88        '''
89        Call the function with the given name and arguments.
90
91        Parameters
92        ----------
93        location
94            From where in the logic program the function was called.
95        name
96            The name of the function.
97        arguments
98            The arguments to the function.
99
100        Returns
101        -------
102        The resulting pool of symbols.
103        '''
104
105    @abstractmethod
106    def callable(self, name: str) -> bool:
107        '''
108        Check there is a function with the given name.
109
110        Parameters
111        ----------
112        name
113            The name of the function.
114
115        Returns
116        -------
117        Whether the function is callable.
118        '''
119
120    def main(self, control: Control) -> None:
121        '''
122        Run the main function.
123
124        This function exisits primarily for internal purposes and does not need
125        to be implemented.
126
127        Parameters
128        ----------
129        control
130            Control object to pass to the main function.
131        '''
132
133class _PythonScript(Script):
134    def execute(self, location, code):
135        exec(code, __main__.__dict__, __main__.__dict__) # pylint: disable=exec-used
136
137    def call(self, location, name, arguments):
138        fun = getattr(__main__, name)
139        return fun(*arguments)
140
141    def callable(self, name):
142        return name in __main__.__dict__ and callable(__main__.__dict__[name])
143
144    def main(self, control):
145        __main__.main(control) # pylint: disable=c-extension-no-member
146
147_PYTHON_SCRIPT = _PythonScript()
148_GLOBAL_SCRIPTS: List[Tuple[Script, Any]] = []
149
150
151@_ffi.def_extern(onerror=_cb_error_top_level, name='pyclingo_script_execute')
152def _pyclingo_script_execute(location, code, data):
153    script: Script = _ffi.from_handle(data)
154    script.execute(_py_location(location), _ffi.string(code).decode())
155    return True
156
157@_ffi.def_extern(onerror=_cb_error_top_level, name='pyclingo_script_call')
158def _pyclingo_script_call(location, name, arguments, size, symbol_callback, symbol_callback_data, data):
159    script: Script = _ffi.from_handle(data)
160    symbol_callback = _ffi.cast('clingo_symbol_callback_t', symbol_callback)
161    arguments = _ffi.cast('clingo_symbol_t*', arguments)
162    py_name = _ffi.string(name).decode()
163    py_args = [Symbol(arguments[i]) for i in range(size)]
164
165    ret = script.call(_py_location(location), py_name, py_args)
166    symbols = list(ret) if isinstance(ret, IterableABC) else [ret]
167
168    c_symbols = _ffi.new('clingo_symbol_t[]', len(symbols))
169    for i, sym in enumerate(symbols):
170        c_symbols[i] = sym._rep # pylint: disable=protected-access
171    _handle_error(symbol_callback(c_symbols, len(symbols), symbol_callback_data))
172    return True
173
174@_ffi.def_extern(onerror=_cb_error_top_level, name='pyclingo_script_callable')
175def _pyclingo_script_callable(name, ret, data):
176    script: Script = _ffi.from_handle(data)
177    py_name = _ffi.string(name).decode()
178    ret[0] = script.callable(py_name)
179    return True
180
181@_ffi.def_extern(onerror=_cb_error_top_level, name='pyclingo_script_main')
182def _pyclingo_script_main(ctl, data):
183    script: Script = _ffi.from_handle(data)
184    script.main(Control(_ffi.cast('clingo_control_t*', ctl)))
185    return True
186
187
188@_ffi.def_extern(onerror=_cb_error_top_level, name='pyclingo_execute')
189def _pyclingo_execute(location, code, data):
190    assert data == _ffi.NULL
191    return _pyclingo_script_execute(_ffi.cast('clingo_location_t*', location),
192                                   code,
193                                   _ffi.new_handle(_PYTHON_SCRIPT))
194
195@_ffi.def_extern(onerror=_cb_error_top_level, name='pyclingo_call')
196def _pyclingo_call(location, name, arguments, size, symbol_callback, symbol_callback_data, data):
197    assert data == _ffi.NULL
198    return _pyclingo_script_call(_ffi.cast('clingo_location_t*', location),
199                                 name,
200                                 arguments, size,
201                                 symbol_callback, symbol_callback_data,
202                                 _ffi.new_handle(_PYTHON_SCRIPT))
203
204@_ffi.def_extern(onerror=_cb_error_top_level, name='pyclingo_callable')
205def _pyclingo_callable(name, ret, data):
206    assert data == _ffi.NULL
207    return _pyclingo_script_callable(name,
208                                     ret,
209                                     _ffi.new_handle(_PYTHON_SCRIPT))
210
211@_ffi.def_extern(onerror=_cb_error_top_level, name='pyclingo_main')
212def _pyclingo_main(ctl, data):
213    assert data == _ffi.NULL
214    return _pyclingo_script_main(_ffi.cast('clingo_control_t*', ctl),
215                                 _ffi.new_handle(_PYTHON_SCRIPT))
216
217
218def register_script(name: str, script: Script, version: str = '1.0.0') -> None:
219    '''
220    Registers a script language which can then be embedded into a logic
221    program.
222
223    Parameters
224    ----------
225    name
226        The name of the script. This name can then be used in the script
227        statement in a logic program.
228    script
229        The class to register.
230    version
231        The version of the script.
232    '''
233    c_version = _c_call('char const*', _lib.clingo_add_string, version.encode())
234    c_script = _ffi.new('clingo_script_t*')
235    c_script[0].execute = _ffi.cast('void*', _lib.pyclingo_script_execute)
236    c_script[0].call = _ffi.cast('void*', _lib.pyclingo_script_call)
237    c_script[0].callable = _ffi.cast('void*', _lib.pyclingo_script_callable)
238    c_script[0].main = _ffi.cast('void*', _lib.pyclingo_script_main)
239    c_script[0].free = _ffi.NULL
240    c_script[0].version = c_version
241    data = _ffi.new_handle(script)
242    _GLOBAL_SCRIPTS.append((script, data))
243    _handle_error(_lib.clingo_register_script(name.encode(), c_script, data))
244
245def enable_python() -> None:
246    '''
247    This function can be called to enable evaluation of Python scripts in logic
248    programs.
249
250    By default evaluation is only enabled in the clingo executable but not in
251    the Python module.
252    '''
253    c_version = _c_call('char const*', _lib.clingo_add_string, python_version().encode())
254    c_script = _ffi.new('clingo_script_t*')
255    c_script[0].execute = _ffi.cast('void*', _lib.pyclingo_execute)
256    c_script[0].call = _ffi.cast('void*', _lib.pyclingo_call)
257    c_script[0].callable = _ffi.cast('void*', _lib.pyclingo_callable)
258    c_script[0].main = _ffi.cast('void*', _lib.pyclingo_main)
259    c_script[0].free = _ffi.NULL
260    c_script[0].version = c_version
261    _handle_error(_lib.clingo_register_script("python".encode(), c_script, _ffi.NULL))
262