Fix review issues in virtualenv PymodeRun integration

diraol · claude · diraol · commit 2aec97cea673 · 2026-05-23T22:39:38.000-03:00
- Set __file__ correctly in subprocess so user code referencing it sees
  the real source path instead of a temp file path
- Rewrite temp file paths in subprocess tracebacks to the real source file
- Add explicit UTF-8 encoding to subprocess Popen call
- Add g:pymode_run_timeout support with TimeoutExpired handling to
  prevent Vim from hanging on non-terminating scripts
- Restore sys.stdout/sys.stderr before returning on SystemExit with
  non-false exit code (pre-existing leak)
- Fix test: use fnameescape() instead of string() for PymodeVirtualenv
  argument (string() wraps in quotes, producing a malformed path)
- Save and restore g:pymode_virtualenv_enabled/g:pymode_virtualenv in
  the test to avoid polluting subsequent tests
- Remove misleading Assert 1 "skip" patterns; drop unreachable else
  branches so cleanup runs unconditionally

Co-Authored-By: Claude Sonnet 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/plugin/pymode.vim b/plugin/pymode.vim
@@ -99,6 +99,9 @@ call pymode#default('g:pymode_run', 1)
 " Key's map for run python code
 call pymode#default('g:pymode_run_bind', '<leader>r')
 
+" Timeout in seconds for :PymodeRun when using a virtualenv interpreter (0 = no limit)
+call pymode#default('g:pymode_run_timeout', 0)
+
 " }}}
 
 " CHECK CODE {{{
diff --git a/pymode/run.py b/pymode/run.py
@@ -28,9 +28,10 @@ def run_code():
     elif encoding.match(lines[1]):
         lines.pop(1)
 
+    real_file = env.var('expand("%:p")')
     python_cmd = __get_virtualenv_python()
     if python_cmd:
-        result = __run_with_external_python(python_cmd, lines)
+        result = __run_with_external_python(python_cmd, lines, real_file)
         if result is not None:
             output, err = result
         else:
@@ -39,7 +40,7 @@ def run_code():
     if not python_cmd:
         context = dict(
             __name__='__main__',
-            __file__=env.var('expand("%:p")'),
+            __file__=real_file,
             input=env.user_input,
             raw_input=env.user_input)
 
@@ -57,6 +58,7 @@ def run_code():
                 # A non-false code indicates abnormal termination.
                 # A false code will be treated as a
                 # successful run, and the error will be hidden from Vim
+                sys.stdout, sys.stderr = stdout_, stderr_
                 env.error("Script exited with code %s" % e.code)
                 return env.stop()
 
@@ -98,8 +100,10 @@ def __get_virtualenv_python():
     return None
 
 
-def __run_with_external_python(python_cmd, lines):
-    source = '\n'.join(lines) + '\n'
+def __run_with_external_python(python_cmd, lines, real_file):
+    # Inject __file__ so user code sees the real source path, not the temp file
+    header = '__file__ = %r\n' % real_file
+    source = header + '\n'.join(lines) + '\n'
     temp_file_path = None
     try:
         with tempfile.NamedTemporaryFile(
@@ -108,20 +112,31 @@ def __run_with_external_python(python_cmd, lines):
             temp_file.write(source)
             temp_file_path = temp_file.name
 
+        timeout = env.var('g:pymode_run_timeout', silence=True, default=0) or None
         process = subprocess.Popen(
             [python_cmd, temp_file_path],
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             cwd=env.curdir,
-            text=True)
-        output, err = process.communicate()
+            text=True,
+            encoding='utf-8')
+        try:
+            output, err = process.communicate(timeout=timeout)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            output, err = process.communicate()
+            err += '\nProcess timed out after %s second(s).' % timeout
     except OSError as exc:
         env.debug('Failed to execute external python', python_cmd, exc)
         return None
     finally:
         if temp_file_path and os.path.exists(temp_file_path):
             os.unlink(temp_file_path)
 
+    # Replace temp path in tracebacks so errors reference the real source file
+    if temp_file_path:
+        err = err.replace(temp_file_path, real_file)
+
     return output, err
 
 
diff --git a/tests/vader/commands.vader b/tests/vader/commands.vader
@@ -94,14 +94,16 @@ Execute (Test PymodeRun uses active virtualenv python):
       call mkdir(temp_dir, 'p')
       let venv_dir = temp_dir . '/.venv'
       let sample_file = temp_dir . '/run_from_venv.py'
+      let saved_venv_enabled = get(g:, 'pymode_virtualenv_enabled', '')
+      let saved_venv = get(g:, 'pymode_virtualenv', 0)
 
       call writefile(['import sys', 'print(sys.executable)'], sample_file)
 
       call system(pycmd . ' -m venv ' . shellescape(venv_dir))
       if v:shell_error == 0
         execute 'edit ' . fnameescape(sample_file)
         let g:pymode_virtualenv = 1
-        execute 'PymodeVirtualenv ' . string(venv_dir)
+        execute 'PymodeVirtualenv ' . fnameescape(venv_dir)
         PymodeRun
 
         let run_buffer = bufnr('__run__')
@@ -114,17 +116,13 @@ Execute (Test PymodeRun uses active virtualenv python):
         else
           Assert 0, 'PymodeRun should create run buffer when using virtualenv'
         endif
-      else
-        Assert 1, 'Unable to create test virtualenv in this environment - test skipped'
       endif
 
+      let g:pymode_virtualenv_enabled = saved_venv_enabled
+      let g:pymode_virtualenv = saved_venv
       call delete(sample_file)
       call delete(temp_dir, 'rf')
-    else
-      Assert 1, 'No python executable available to create virtualenv - test skipped'
     endif
-  else
-    Assert 1, 'PymodeRun/PymodeVirtualenv command not available - test skipped'
   endif
 
 # Test PymodeLint command
