From 9b314777e0000f3e94c33aab1ef19ab52ddc242d Mon Sep 17 00:00:00 2001
From: Laurens Valk <laurens@pybricks.com>
Date: Tue, 29 Nov 2022 13:00:27 +0100
Subject: [PATCH] micropython: Document heap_lock, heap_unlock.

Fixes https://github.com/pybricks/pybricks-api/issues/116
---
 doc/main/micropython/micropython.rst | 12 ++++++++----
 jedi/tests/test_complete_import.py   |  2 ++
 src/micropython/__init__.py          | 22 +++++++++++++++++++++-
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/doc/main/micropython/micropython.rst b/doc/main/micropython/micropython.rst
index 73e4497..71868bf 100644
--- a/doc/main/micropython/micropython.rst
+++ b/doc/main/micropython/micropython.rst
@@ -6,16 +6,20 @@
 
 .. autofunction:: micropython.const
 
-.. autofunction:: micropython.opt_level
+.. autofunction:: micropython.heap_lock
+
+.. autofunction:: micropython.heap_unlock
+
+.. autofunction:: micropython.kbd_intr
 
 .. autofunction:: micropython.mem_info
 
+.. autofunction:: micropython.opt_level
+
 .. autofunction:: micropython.qstr_info
-
+    
 .. autofunction:: micropython.stack_use
 
-.. autofunction:: micropython.kbd_intr
-
 
 Examples
 ---------------------
diff --git a/jedi/tests/test_complete_import.py b/jedi/tests/test_complete_import.py
index ee874e6..7e85928 100644
--- a/jedi/tests/test_complete_import.py
+++ b/jedi/tests/test_complete_import.py
@@ -149,6 +149,8 @@ def test_from_micropython_import():
     completions: list[CompletionItem] = json.loads(complete(code, 1, len(code) + 1))
     assert [c["insertText"] for c in completions] == [
         "const",
+        "heap_lock",
+        "heap_unlock",
         "kbd_intr",
         "mem_info",
         "opt_level",
diff --git a/src/micropython/__init__.py b/src/micropython/__init__.py
index 39577f1..2b6ae2d 100644
--- a/src/micropython/__init__.py
+++ b/src/micropython/__init__.py
@@ -108,7 +108,27 @@ def stack_use() -> int:
     """
 
 
-# REVISIT: Skipping heap lock funcs for now
+def heap_lock() -> None:
+    """
+    heap_lock()
+
+    Locks the heap. When locked, no memory allocation can occur. A
+    ``MemoryError`` will be raised if any heap allocation is attempted.
+    """
+
+
+def heap_unlock() -> int:
+    """
+    heap_unlock() -> int
+
+    Unlocks the heap. Memory allocation is now allowed again.
+
+    If :func:`heap_lock()` was called multiple times, :func:`heap_unlock()`
+    must be called the same number of times to make the heap available again.
+
+    Returns:
+        The lock depth after unlocking. It is ``0`` once it is unlocked.
+    """
 
 
 def kbd_intr(chr: int) -> None: