From 6280a07f56d0a983769ac60be762875376324d1c Mon Sep 17 00:00:00 2001
From: Wu Zhenyu <wuzy01@qq.com>
Date: Sun, 28 Aug 2022 12:45:00 +0800
Subject: [PATCH] Add completes for shells by shtab

ptpython --print-completion bash | sudo tee /usr/share/bash-completion/completions/ptpython
ptpython --print-completion zsh | sed 's/compdef ptpython/compdef -P pt(i|)python[0-9.]#/' | sudo tee /usr/share/zsh/site-functions/_ptpython  # wait <https://github.com/iterative/shtab/issues/87>
ptpython --print-completion tcsh | sudo tee /etc/profile.d/ptpython.completion.csh
---
 ptpython/entry_points/run_ptpython.py | 37 +++++++++++++++++++++++----
 setup.py                              |  1 +
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/ptpython/entry_points/run_ptpython.py b/ptpython/entry_points/run_ptpython.py
index 5ebe2b9..e52dcb4 100644
--- a/ptpython/entry_points/run_ptpython.py
+++ b/ptpython/entry_points/run_ptpython.py
@@ -28,6 +28,7 @@
 from textwrap import dedent
 from typing import Tuple
 
+import shtab
 import appdirs
 from prompt_toolkit.formatted_text import HTML
 from prompt_toolkit.shortcuts import print_formatted_text
@@ -41,6 +42,27 @@
 
 
 __all__ = ["create_parser", "get_config_and_history_file", "run"]
+# https://github.com/iterative/shtab/blob/master/examples/customcomplete.py#L11-L22
+PY_FILE = {
+    "bash": "_shtab_greeter_compgen_py_file",
+    "zsh": "_files -g *.py",
+    "tcsh": "f:*.txt",
+}
+PREAMBLE = {
+    "bash": """\
+# $1=COMP_WORDS[1]
+_shtab_greeter_compgen_py_file() {
+  compgen -d -- $1  # recurse into subdirs
+  compgen -f -X '!*?.txt' -- $1
+}
+""",
+    "zsh": """\
+_script_args() {
+  _arguments -S -s '(-)1:script_args:_files -g *.py' '*: :_files'
+}
+""",
+}
+SCRIPT_ARGS = {"zsh": "_script_args"}
 
 
 class _Parser(argparse.ArgumentParser):
@@ -59,6 +81,7 @@ def print_help(self):
 
 def create_parser() -> _Parser:
     parser = _Parser(description="ptpython: Interactive Python shell.")
+    shtab.add_argument_to(parser, preamble=PREAMBLE)
     parser.add_argument("--vi", action="store_true", help="Enable Vi key bindings")
     parser.add_argument(
         "-i",
@@ -70,23 +93,27 @@ def create_parser() -> _Parser:
         "--light-bg",
         action="store_true",
         help="Run on a light background (use dark colors for text).",
-    ),
+    )
     parser.add_argument(
         "--dark-bg",
         action="store_true",
         help="Run on a dark background (use light colors for text).",
-    ),
+    )
     parser.add_argument(
         "--config-file", type=str, help="Location of configuration file."
-    )
-    parser.add_argument("--history-file", type=str, help="Location of history file.")
+    ).complete = PY_FILE
+    parser.add_argument(
+        "--history-file", type=str, help="Location of history file."
+    ).complete = shtab.FILE
     parser.add_argument(
         "-V",
         "--version",
         action="version",
         version=metadata.version("ptpython"),  # type: ignore
     )
-    parser.add_argument("args", nargs="*", help="Script and arguments")
+    parser.add_argument(
+        "args", nargs=argparse.REMAINDER, help="Script and arguments"
+    ).complete = SCRIPT_ARGS
     return parser
 
 
diff --git a/setup.py b/setup.py
index a8214f2..7902bac 100644
--- a/setup.py
+++ b/setup.py
@@ -18,6 +18,7 @@
     packages=find_packages("."),
     package_data={"ptpython": ["py.typed"]},
     install_requires=[
+        "shtab",
         "appdirs",
         "importlib_metadata;python_version<'3.8'",
         "jedi>=0.16.0",