Skip to content

Commit b33d048

Browse files
CarreauCarreau
authored
execution count more atomic, run_cell_async reenterable (#15088)
This pull request makes `run_cell_async` able to re-enter itself without messing the history causing a new session. This solves issue #15087 The changes are to increment `execution_count` of `InteractiveShell` _before_ running the user code in a cell, and some compatible changes so that prompts and outputs get the correct `execution_count`. With this change it's possible to modify frontends so that they can run multiple cells concurrently. I've written a small one-file package on top of ipykernel that shows this nicely and works well https://github.com/drorspei/bgipykernel
2 parents 263ebaa + aaa5a45 commit b33d048

5 files changed

Lines changed: 39 additions & 18 deletions

File tree

‎IPython/core/displayhook.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ def __init__(self, shell=None, cache_size=1000, **kwargs):
6363

6464
@property
6565
def prompt_count(self):
66-
return self.shell.execution_count
66+
return self.shell.execution_count - 1
6767

6868
#-------------------------------------------------------------------------
6969
# Methods used in __call__. Override these methods to modify the behavior
@@ -127,7 +127,7 @@ def write_output_prompt(self):
127127
"""
128128
# Use write, not print which adds an extra space.
129129
sys.stdout.write(self.shell.separate_out)
130-
outprompt = 'Out[{}]: '.format(self.shell.execution_count)
130+
outprompt = 'Out[{}]: '.format(self.shell.execution_count - 1)
131131
if self.do_full_cache:
132132
sys.stdout.write(outprompt)
133133

‎IPython/core/displaypub.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ def publish(
149149

150150
outputs = self.shell.history_manager.outputs
151151

152-
target_execution_count = self.shell.execution_count
152+
target_execution_count = self.shell.execution_count - 1
153153
if self._in_post_execute:
154154
# We're in post_execute, so this is likely a matplotlib flush
155155
# Use execution_count - 1 to associate with the cell that created the plot

‎IPython/core/interactiveshell.py‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3046,6 +3046,7 @@ def _tee(self, channel: Literal["stdout", "stderr"]):
30463046
"""
30473047
stream = getattr(sys, channel)
30483048
original_write = stream.write
3049+
execution_count = self.execution_count
30493050

30503051
def write(data, *args, **kwargs):
30513052
"""Write data to both the original destination and the capture dictionary."""
@@ -3060,7 +3061,6 @@ def write(data, *args, **kwargs):
30603061
return result
30613062
if not data:
30623063
return result
3063-
execution_count = self.execution_count
30643064
output_stream = None
30653065
outputs_by_counter = self.history_manager.outputs
30663066
output_type = "out_stream" if channel == "stdout" else "err_stream"
@@ -3293,17 +3293,18 @@ async def run_cell_async(
32933293
if silent:
32943294
store_history = False
32953295

3296+
execution_count = result.execution_count = self.execution_count
3297+
32963298
if store_history:
3297-
result.execution_count = self.execution_count
3299+
self.execution_count += 1
32983300

32993301
def error_before_exec(value):
33003302
if store_history:
33013303
if self.history_manager:
33023304
# Store formatted traceback and error details
33033305
self.history_manager.exceptions[
3304-
self.execution_count
3306+
execution_count
33053307
] = self._format_exception_for_storage(value)
3306-
self.execution_count += 1
33073308
result.error_before_exec = value
33083309
self.last_execution_succeeded = False
33093310
self.last_execution_result = result
@@ -3348,15 +3349,13 @@ def error_before_exec(value):
33483349
# Store raw and processed history
33493350
if store_history:
33503351
assert self.history_manager is not None
3351-
self.history_manager.store_inputs(self.execution_count, cell, raw_cell)
3352+
self.history_manager.store_inputs(execution_count, cell, raw_cell)
33523353
if not silent:
33533354
self.logger.log(cell, raw_cell)
33543355

33553356
# Display the exception if input processing failed.
33563357
if preprocessing_exc_tuple is not None:
33573358
self.showtraceback(preprocessing_exc_tuple)
3358-
if store_history:
3359-
self.execution_count += 1
33603359
return error_before_exec(preprocessing_exc_tuple[1])
33613360

33623361
# Our own compiler remembers the __future__ environment. If we want to
@@ -3365,7 +3364,7 @@ def error_before_exec(value):
33653364
compiler = self.compile if shell_futures else self.compiler_class()
33663365

33673366
with self.builtin_trap:
3368-
cell_name = compiler.cache(cell, self.execution_count, raw_code=raw_cell)
3367+
cell_name = compiler.cache(cell, execution_count, raw_code=raw_cell)
33693368

33703369
with self.display_trap:
33713370
# Compile to bytecode
@@ -3412,17 +3411,13 @@ def error_before_exec(value):
34123411
assert self.history_manager is not None
34133412
# Write output to the database. Does nothing unless
34143413
# history output logging is enabled.
3415-
self.history_manager.store_output(self.execution_count)
3416-
exec_count = self.execution_count
3414+
self.history_manager.store_output(execution_count)
34173415
if result.error_in_exec:
34183416
# Store formatted traceback and error details
34193417
self.history_manager.exceptions[
3420-
exec_count
3418+
execution_count
34213419
] = self._format_exception_for_storage(result.error_in_exec)
34223420

3423-
# Each cell is a *single* input, regardless of how many lines it has
3424-
self.execution_count += 1
3425-
34263421
return result
34273422

34283423
def _format_exception_for_storage(

‎IPython/terminal/prompts.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ def rewrite_prompt_tokens(self):
8787
def out_prompt_tokens(self):
8888
return [
8989
(Token.OutPrompt, 'Out['),
90-
(Token.OutPromptNum, str(self.shell.execution_count)),
90+
(Token.OutPromptNum, str(self.shell.execution_count - 1)),
9191
(Token.OutPrompt, ']: '),
9292
]
9393

‎tests/test_history.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -375,3 +375,29 @@ def test_get_tail_session_awareness(hmmax3):
375375
hm2.db.close()
376376
if ha:
377377
ha.db.close()
378+
379+
380+
def test_calling_run_cell(hmmax2):
381+
ip = get_ipython()
382+
with TemporaryDirectory() as tmpdir:
383+
tmp_path = Path(tmpdir)
384+
hist_manager_ori = ip.history_manager
385+
hist_file = tmp_path / "history_test_history1.sqlite"
386+
try:
387+
ip.history_manager = HistoryManager(shell=ip, hist_file=hist_file)
388+
import time
389+
390+
session_number = ip.history_manager.session_number
391+
ip.run_cell(raw_cell="get_ipython().run_cell(raw_cell='1', store_history=True)", store_history=True)
392+
while ip.history_manager.db_input_cache:
393+
time.sleep(0)
394+
new_session_number = ip.history_manager.session_number
395+
finally:
396+
# Ensure saving thread is shut down before we try to clean up the files
397+
ip.history_manager.end_session()
398+
# Forcibly close database rather than relying on garbage collection
399+
ip.history_manager.save_thread.stop()
400+
ip.history_manager.db.close()
401+
402+
ip.history_manager = hist_manager_ori
403+
assert session_number == new_session_number, ValueError(f"{session_number} != {new_session_number}")

0 commit comments

Comments
 (0)