Update REPL based on asyncio.__main__ in CPython 3.14

README.md

@@ -55,6 +55,8 @@ Combine these with `-i` to open a REPL after running the file specified on the c
 
 See `import-expression --help` for more details.
 
+The REPL requires Python 3.10+.
+
 ### Running a file
 
 Run `import-expression <filename.py>`.

import_expression/__main__.py

@@ -1,4 +1,5 @@
 # Copyright © io mintz <io@mintz.cc>
+# Copyright © Thanos <111999343+Sachaa-Thanasius@users.noreply.github.com>
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the “Software”),
@@ -30,36 +31,32 @@
 import codeop
 import concurrent.futures
 import contextlib
+import contextvars
 import importlib
 import inspect
 import os.path
 import rlcompleter
 import sys
 import traceback
 import threading
+import tokenize
 import types
 import warnings
 from asyncio import futures
-from codeop import PyCF_DONT_IMPLY_DEDENT, PyCF_ALLOW_INCOMPLETE_INPUT
+
+try:
+	from codeop import PyCF_DONT_IMPLY_DEDENT, PyCF_ALLOW_INCOMPLETE_INPUT
+except ImportError:
+	raise RuntimeError('The import-expression interactive REPL is only supported on Python 3.10+.')
 
 import import_expression
 from import_expression import constants
 
-if os.path.basename(sys.argv[0]) == 'import_expression':
-	import warnings
-	warnings.warn(UserWarning(
-		'The import_expression alias is deprecated, and will be removed in v2.0. '
-		'Please use import-expression (with a hyphen) instead.'
-	))
-
 features = [getattr(__future__, fname) for fname in __future__.all_feature_names]
 
-try:
-	from ast import PyCF_ALLOW_TOP_LEVEL_AWAIT
-except ImportError:
-	SUPPORTS_ASYNCIO_REPL = False
-else:
-	SUPPORTS_ASYNCIO_REPL = True
+from ast import PyCF_ALLOW_TOP_LEVEL_AWAIT
+
+return_code = 0
 
 class ImportExpressionCommandCompiler(codeop.CommandCompiler):
 	def __init__(self):
@@ -75,12 +72,14 @@ class ImportExpressionCompile:
 	def __init__(self):
 		self.flags = PyCF_DONT_IMPLY_DEDENT | PyCF_ALLOW_INCOMPLETE_INPUT
 
-	def __call__(self, source, filename, symbol, **kwargs):
-		flags = self.flags
+	def __call__(self, source, filename, symbol, flags=0, **kwargs):
+		flags |= self.flags
 		if kwargs.get('incomplete_input', True) is False:
 			flags &= ~PyCF_DONT_IMPLY_DEDENT
 			flags &= ~PyCF_ALLOW_INCOMPLETE_INPUT
 		codeob = import_expression.compile(source, filename, symbol, flags, True)
+		if flags & ast.PyCF_ONLY_AST:
+			return codeob   # this is an ast.Module in this case
 		for feature in features:
 			if codeob.co_flags & feature.compiler_flag:
 				self.flags |= feature.compiler_flag
@@ -100,11 +99,14 @@ def __init__(self, locals, loop):
 		self.compile.compiler.flags |= PyCF_ALLOW_TOP_LEVEL_AWAIT
 
 		self.loop = loop
+		self.context = contextvars.copy_context()
 
 	def runcode(self, code):
+		global return_code
 		future = concurrent.futures.Future()
 
 		def callback():
+			global return_code
 			global repl_future
 			global repl_future_interrupted
 
@@ -114,8 +116,10 @@ def callback():
 			func = types.FunctionType(code, self.locals)
 			try:
 				coro = func()
-			except SystemExit:
-				raise
+			except SystemExit as se:
+				return_code = se.code
+				self.loop.stop()
+				return
 			except KeyboardInterrupt as ex:
 				repl_future_interrupted = True
 				future.set_exception(ex)
@@ -129,30 +133,53 @@ def callback():
 				return
 
 			try:
-				repl_future = self.loop.create_task(coro)
+				repl_future = self.loop.create_task(coro, context=self.context)
 				futures._chain_future(repl_future, future)
 			except BaseException as exc:
 				future.set_exception(exc)
 
-		self.loop.call_soon_threadsafe(callback)
+		self.loop.call_soon_threadsafe(callback, context=self.context)
 
 		try:
 			return future.result()
-		except SystemExit:
-			raise
+		except SystemExit as se:
+			return_code = se.code
+			self.loop.stop()
+			return
 		except BaseException:
 			if repl_future_interrupted:
 				self.write("\nKeyboardInterrupt\n")
 			else:
 				self.showtraceback()
 
 class REPLThread(threading.Thread):
-	def __init__(self, interact_kwargs):
-		self.interact_kwargs = interact_kwargs
+	def __init__(self, interact_kwargs, prelude_path: str | None):
 		super().__init__()
+		self.interact_kwargs = interact_kwargs
+		self.prelude_path = prelude_path
 
 	def run(self):
 		try:
+			if startup_path := os.getenv("PYTHONSTARTUP"):
+				with tokenize.open(startup_path) as f:
+					startup_code = compile(f.read(), startup_path, "exec")
+					exec(startup_code, console.locals)
+
+			if self.prelude_path is not None:
+				# ensure loop is running
+				asyncio.run_coroutine_threadsafe(asyncio.sleep(0), loop).result()
+
+				with tokenize.open(self.prelude_path) as f:
+					prelude_source = f.read()
+
+				prelude_code = import_expression.compile(
+					prelude_source,
+					self.prelude_path,
+					'exec',
+					PyCF_ALLOW_TOP_LEVEL_AWAIT
+				)
+				console.runcode(prelude_code)
+
 			console.interact(**self.interact_kwargs)
 		finally:
 			warnings.filterwarnings(
@@ -189,34 +216,37 @@ def attr_matches(self, text):
 		self.namespace = old_namespace
 		return res
 
-def asyncio_main(repl_locals, interact_kwargs):
+def asyncio_main(repl_locals, interact_kwargs, prelude_path: str | None):
 	global console
 	global loop
 	global repl_future
 	global repl_future_interrupted
 
-	loop = asyncio.get_event_loop()
+	loop = asyncio.new_event_loop()
+	asyncio.set_event_loop(loop)
 
 	console = ImportExpressionAsyncIOInteractiveConsole(repl_locals, loop)
 
 	repl_future = None
 	repl_future_interrupted = False
 
-	repl_thread = REPLThread(interact_kwargs)
+	repl_thread = REPLThread(interact_kwargs, prelude_path)
 	repl_thread.daemon = True
 	repl_thread.start()
 
 	while True:
 		try:
 			loop.run_forever()
 		except KeyboardInterrupt:
+			repl_future_interrupted = True
 			if repl_future and not repl_future.done():
 				repl_future.cancel()
-				repl_future_interrupted = True
 			continue
 		else:
 			break
 
+	return return_code
+
 def parse_args():
 	import argparse
 
@@ -264,32 +294,24 @@ def main():
 
 	repl_locals = {
 		key: globals()[key] for key in [
-			'__name__', '__package__',
-			'__loader__', '__spec__',
-			'__builtins__', '__file__'
+			'__package__', '__loader__',
+			'__spec__', '__builtins__',
+			'__file__',
 		]
 		if key in globals()
 	}
 
-	args = parse_args()
+	repl_locals['__name__'] = '__main__'
 
-	if args.asyncio and not SUPPORTS_ASYNCIO_REPL:
-		print('Python3.8+ required for the AsyncIO REPL.', file=sys.stderr)
-		sys.exit(2)
+	args = parse_args()
 
-	prelude = None
 	if args.filename:
-		with open(args.filename) as f:
-			flags = 0
-			if args.asyncio:
-				flags |= PyCF_ALLOW_TOP_LEVEL_AWAIT
-			prelude = import_expression.compile(f.read(), flags=flags)
-		if args.asyncio:
-			# we need a new loop because using asyncio.run here breaks the console
-			loop = asyncio.new_event_loop()
-			loop.run_until_complete(eval(prelude, repl_locals))
-		else:
-			import_expression.exec(prelude, globals=repl_locals)
+		if not args.asyncio:
+			with open(args.filename) as f:
+				prelude_source = f.read()
+			codeobj = import_expression.compile(prelude_source, args.filename)
+			import_expression.exec(codeobj, globals=repl_locals)
+
 		if not args.interactive:
 			sys.exit(0)
 
@@ -298,8 +320,7 @@ def main():
 	interact_kwargs = dict(banner='' if args.quiet else None, exitmsg='' if args.quiet else None)
 
 	if args.asyncio:
-		asyncio_main(repl_locals, interact_kwargs)
-		sys.exit(0)
+		sys.exit(asyncio_main(repl_locals, interact_kwargs, args.filename or None))
 
 	ImportExpressionInteractiveConsole(repl_locals).interact(**interact_kwargs)