Advanced tempfile

Introduced a new tempfile argument to Buffer which allows the user
to provide a partial path in order to enable more advanced syntax
highlighting (e.g. git commit message highlighting).

Author: Ken Odegard

prompt_toolkit/buffer.py

@@ -6,6 +6,7 @@
 import os
 import re
 import shlex
+import shutil
 import subprocess
 import tempfile
 from asyncio import Task, ensure_future
@@ -167,6 +168,10 @@ class Buffer:
         "open in editor" function. For a Python REPL, this would be ".py", so
         that the editor knows the syntax highlighting to use. This can also be
         a callable that returns a string.
+    :param tempfile: For more advanced tempfile situations where you need
+        control over the subdirectories and filename. For a Git Commit Message,
+        this would be ".git/COMMIT_EDITMSG", so that the editor knows the syntax
+        highlighting to use. This can also be a callable that returns a string.
     :param name: Name for this buffer. E.g. DEFAULT_BUFFER. This is mostly
         useful for key bindings where we sometimes prefer to refer to a buffer
         by their name instead of by reference.
@@ -213,6 +218,7 @@ def __init__(self,
                  history: Optional[History] = None,
                  validator: Optional[Validator] = None,
                  tempfile_suffix: Union[str, Callable[[], str]] = '',
+                 tempfile: Union[str, Callable[[], str]] = '',
                  name: str = '',
                  complete_while_typing: FilterOrBool = False,
                  validate_while_typing: FilterOrBool = False,
@@ -238,6 +244,7 @@ def __init__(self,
         self.auto_suggest = auto_suggest
         self.validator = validator
         self.tempfile_suffix = tempfile_suffix
+        self.tempfile = tempfile
         self.name = name
         self.accept_handler = accept_handler
 
@@ -1398,7 +1405,30 @@ def open_in_editor(self, validate_and_handle: bool = False) -> 'asyncio.Task[Non
             raise EditReadOnlyBuffer()
 
         # Write to temporary file
-        descriptor, filename = tempfile.mkstemp(to_str(self.tempfile_suffix))
+        if self.tempfile:
+            # Try to make according to tempfile logic.
+            head, tail = os.path.split(to_str(self.tempfile))
+            if os.path.isabs(head):
+                head = head[1:]
+
+            dirpath = tempfile.mkdtemp()
+            remove = dirpath
+            if head:
+                dirpath = os.path.join(dirpath, head)
+            # Assume there is no issue creating dirs in this temp dir.
+            os.makedirs(dirpath)
+
+            # Open the filename of interest.
+            filename = os.path.join(dirpath, tail)
+            descriptor = os.open(filename, os.O_WRONLY|os.O_CREAT)
+        else:
+            # Fallback to tempfile_suffix logic.
+            suffix = None
+            if self.tempfile_suffix:
+                suffix = to_str(self.tempfile_suffix)
+            descriptor, filename = tempfile.mkstemp(suffix)
+            remove = filename
+
         os.write(descriptor, self.text.encode('utf-8'))
         os.close(descriptor)
 
@@ -1430,8 +1460,8 @@ async def run() -> None:
                         self.validate_and_handle()
 
             finally:
-                # Clean up temp file.
-                os.remove(filename)
+                # Clean up temp dir/file.
+                shutil.rmtree(remove)
 
         return get_app().create_background_task(run())
 

prompt_toolkit/shortcuts/prompt.py

@@ -334,7 +334,8 @@ class PromptSession(Generic[_T]):
         'validate_while_typing', 'complete_style', 'mouse_support',
         'auto_suggest', 'clipboard', 'validator', 'refresh_interval',
         'input_processors', 'enable_system_prompt', 'enable_suspend',
-        'enable_open_in_editor', 'reserve_space_for_menu', 'tempfile_suffix')
+        'enable_open_in_editor', 'reserve_space_for_menu', 'tempfile_suffix',
+        'tempfile')
 
     def __init__(
             self,
@@ -373,7 +374,8 @@ def __init__(
             input_processors: Optional[List[Processor]] = None,
             key_bindings: Optional[KeyBindingsBase] = None,
             erase_when_done: bool = False,
-            tempfile_suffix: str = '.txt',
+            tempfile_suffix: Optional[Union[str, Callable[[], str]]] = '.txt',
+            tempfile: Optional[Union[str, Callable[[], str]]] = None,
 
             refresh_interval: float = 0,
 
@@ -425,6 +427,7 @@ def __init__(
         self.enable_open_in_editor = enable_open_in_editor
         self.reserve_space_for_menu = reserve_space_for_menu
         self.tempfile_suffix = tempfile_suffix
+        self.tempfile = tempfile
 
         # Create buffers, layout and Application.
         self.history = history
@@ -480,7 +483,8 @@ def accept(buff: Buffer) -> bool:
             history=self.history,
             auto_suggest=DynamicAutoSuggest(lambda: self.auto_suggest),
             accept_handler=accept,
-            tempfile_suffix=lambda: self.tempfile_suffix)
+            tempfile_suffix=lambda: self.tempfile_suffix,
+            tempfile=lambda: self.tempfile)
 
     def _create_search_buffer(self) -> Buffer:
         return Buffer(name=SEARCH_BUFFER)
@@ -795,7 +799,8 @@ def prompt(
             enable_system_prompt: Optional[FilterOrBool] = None,
             enable_suspend: Optional[FilterOrBool] = None,
             enable_open_in_editor: Optional[FilterOrBool] = None,
-            tempfile_suffix: Optional[str] = None,
+            tempfile_suffix: Optional[Union[str, Callable[[], str]]] = None,
+            tempfile: Optional[Union[str, Callable[[], str]]] = None,
 
             # Following arguments are specific to the current `prompt()` call.
             default: Union[str, Document] = '',
@@ -910,6 +915,8 @@ class itself. For these, passing in ``None`` will keep the current
             self.enable_open_in_editor = enable_open_in_editor
         if tempfile_suffix is not None:
             self.tempfile_suffix = tempfile_suffix
+        if tempfile is not None:
+            self.tempfile = tempfile
 
         self._add_pre_run_callables(pre_run, accept_default)
         self.default_buffer.reset(default if isinstance(default, Document) else Document(default))
@@ -958,7 +965,8 @@ async def prompt_async(
             enable_system_prompt: Optional[FilterOrBool] = None,
             enable_suspend: Optional[FilterOrBool] = None,
             enable_open_in_editor: Optional[FilterOrBool] = None,
-            tempfile_suffix: Optional[str] = None,
+            tempfile_suffix: Optional[Union[str, Callable[[], str]]] = None,
+            tempfile: Optional[Union[str, Callable[[], str]]] = None,
 
             # Following arguments are specific to the current `prompt()` call.
             default: Union[str, Document] = '',
@@ -1033,6 +1041,8 @@ async def prompt_async(
             self.enable_open_in_editor = enable_open_in_editor
         if tempfile_suffix is not None:
             self.tempfile_suffix = tempfile_suffix
+        if tempfile is not None:
+            self.tempfile = tempfile
 
         self._add_pre_run_callables(pre_run, accept_default)
         self.default_buffer.reset(default if isinstance(default, Document) else Document(default))
@@ -1207,7 +1217,8 @@ def prompt(
         enable_system_prompt: Optional[FilterOrBool] = None,
         enable_suspend: Optional[FilterOrBool] = None,
         enable_open_in_editor: Optional[FilterOrBool] = None,
-        tempfile_suffix: Optional[str] = None,
+        tempfile_suffix: Optional[Union[str, Callable[[], str]]] = None,
+        tempfile: Optional[Union[str, Callable[[], str]]] = None,
 
         # Following arguments are specific to the current `prompt()` call.
         default: str = '',
@@ -1255,6 +1266,7 @@ def prompt(
         enable_suspend=enable_suspend,
         enable_open_in_editor=enable_open_in_editor,
         tempfile_suffix=tempfile_suffix,
+        tempfile=tempfile,
         default=default,
         accept_default=accept_default,
         pre_run=pre_run)