From 4161b869a3a0b5282b2f880696f67b2a9fb7a6de 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/_shtab.py                    |  8 +++++
 ptpython/entry_points/run_ptpython.py | 42 +++++++++++++++++++++++----
 setup.py                              |  2 ++
 3 files changed, 46 insertions(+), 6 deletions(-)
 create mode 100644 ptpython/_shtab.py

diff --git a/ptpython/_shtab.py b/ptpython/_shtab.py
new file mode 100644
index 0000000..cbdce10
--- /dev/null
+++ b/ptpython/_shtab.py
@@ -0,0 +1,8 @@
+FILE = None
+DIRECTORY = DIR = None
+
+
+def add_argument_to(parser, *args, **kwargs):
+    from argparse import Action
+    Action.complete = None
+    return parser
diff --git a/ptpython/entry_points/run_ptpython.py b/ptpython/entry_points/run_ptpython.py
index 5ebe2b9..9fd3ffa 100644
--- a/ptpython/entry_points/run_ptpython.py
+++ b/ptpython/entry_points/run_ptpython.py
@@ -28,6 +28,10 @@
 from textwrap import dedent
 from typing import Tuple
 
+try:
+    import shtab
+except ImportError:
+    from .. import _shtab as shtab
 import appdirs
 from prompt_toolkit.formatted_text import HTML
 from prompt_toolkit.shortcuts import print_formatted_text
@@ -41,6 +45,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:*.py",
+}
+PREAMBLE = {
+    "bash": """\
+# $1=COMP_WORDS[1]
+_shtab_greeter_compgen_py_file() {
+  compgen -d -- $1  # recurse into subdirs
+  compgen -f -X '!*?.py' -- $1
+}
+""",
+    "zsh": """\
+_script_args() {
+  _arguments -S -s '(-)1:script_args:_files -g "*.py"' '*: :_files'
+}
+""",
+}
+SCRIPT_ARGS = {"zsh": "_script_args"}
 
 
 class _Parser(argparse.ArgumentParser):
@@ -58,7 +83,8 @@ def print_help(self):
 
 
 def create_parser() -> _Parser:
-    parser = _Parser(description="ptpython: Interactive Python shell.")
+    parser = _Parser("ptpython", 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 +96,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..080f8a5 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",
@@ -50,5 +51,6 @@
     extras_require={
         "ptipython": ["ipython"],  # For ptipython, we need to have IPython
         "all": ["black"],  # Black not always possible on PyPy
+        "completion": ["shtab"],
     },
 )