From 1bb9c9752c028723d46c06637a0147e73124674c Mon Sep 17 00:00:00 2001
From: MrNaif2018 <chuff184@gmail.com>
Date: Wed, 7 Jul 2021 15:02:04 +0300
Subject: [PATCH] Export only marked functions to context

---
 bitccl/compiler.py     | 5 +++--
 bitccl/functions.py    | 8 +++++++-
 bitccl/utils.py        | 6 ++++++
 tests/test_compiler.py | 1 +
 4 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/bitccl/compiler.py b/bitccl/compiler.py
index 5fedb6d..61ccfd0 100644
--- a/bitccl/compiler.py
+++ b/bitccl/compiler.py
@@ -16,8 +16,9 @@
 from bitccl.exceptions import CompilationRestrictedError
 from bitccl.utils import no_imports_importer
 
-# TODO: currently it exports more than needed
-functions = {name: func for (name, func) in inspect.getmembers(functions_module, inspect.isfunction)}
+functions = {
+    name: func for (name, func) in inspect.getmembers(functions_module, inspect.isfunction) if hasattr(func, "bitccl")
+}
 events = {
     name: event
     for (name, event) in inspect.getmembers(events_module, inspect.isclass)
diff --git a/bitccl/functions.py b/bitccl/functions.py
index 8c62eee..7325cb0 100644
--- a/bitccl/functions.py
+++ b/bitccl/functions.py
@@ -7,17 +7,19 @@
 
 from bitccl.state import config as config_ctx
 from bitccl.state import event_listeners
-from bitccl.utils import call_universal, prepare_event, silent_debug, time_limit
+from bitccl.utils import call_universal, function, prepare_event, silent_debug, time_limit
 
 PASSWORD_ALPHABET = string.ascii_uppercase + string.ascii_lowercase + string.digits + "-_"
 SECURE_PASSWORD_LENGTH = 43
 
 
+@function
 def add_event_listener(event, func):
     event = prepare_event(event)
     event_listeners[event].append(func)
 
 
+@function
 def on(event):
     def wrapper(func):
         add_event_listener(event, func)
@@ -25,6 +27,7 @@ def wrapper(func):
     return wrapper
 
 
+@function
 def dispatch_event(event, *args, **kwargs):
     event = prepare_event(event)
     for listener in event_listeners[event]:
@@ -36,6 +39,7 @@ def dispatch_event(event, *args, **kwargs):
             pass
 
 
+@function
 def template(name, data={}):
     try:
         with open(f"templates/{name}.j2") as f:
@@ -46,6 +50,7 @@ def template(name, data={}):
         return ""
 
 
+@function
 def send_email(to, subject, text):  # pragma: no cover
     try:
         config = config_ctx.get()
@@ -66,5 +71,6 @@ def send_email(to, subject, text):  # pragma: no cover
         return False
 
 
+@function
 def password(length=SECURE_PASSWORD_LENGTH):
     return "".join(secrets.choice(PASSWORD_ALPHABET) for _ in range(length))
diff --git a/bitccl/utils.py b/bitccl/utils.py
index 2111ff6..5908a46 100644
--- a/bitccl/utils.py
+++ b/bitccl/utils.py
@@ -69,3 +69,9 @@ def call_universal(func, *args, **kwargs):
     if inspect.isawaitable(result):
         result = asyncio.get_event_loop().run_until_complete(result)
     return result
+
+
+def function(f):
+    """Use function decorator to mark functions to be exported to BitCCL execution context"""
+    f.bitccl = True
+    return f
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index 023d7cf..09eb962 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -99,6 +99,7 @@ def test_disable_imports():
 
 @pytest.mark.parametrize("func", functions.keys())
 def test_disallowed_imports(func):
+    assert hasattr(functions[func], "bitccl")
     assert "Imports disabled" in run(f"from functions import {func}")