From 70098556e85b5a9fbd5c48fdb46f60a646089bea Mon Sep 17 00:00:00 2001
From: Richard Chien <richardchienthebest@gmail.com>
Date: Sat, 22 Dec 2018 23:06:56 +0800
Subject: [PATCH] Improve ArgumentParser class

---
 none/argparse.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/none/argparse.py b/none/argparse.py
index a5a6a024..040feb79 100644
--- a/none/argparse.py
+++ b/none/argparse.py
@@ -1,5 +1,7 @@
 from argparse import *
 
+from .command import CommandSession
+
 
 class ParserExit(RuntimeError):
     def __init__(self, status=0, message=None):
@@ -13,9 +15,30 @@ class ArgumentParser(ArgumentParser):
     standard I/O.
     """
 
+    def __init__(self, *args, **kwargs):
+        self.session = kwargs.pop('session', None)
+        super().__init__(*args, **kwargs)
+
     def _print_message(self, *args, **kwargs):
         # do nothing
         pass
 
     def exit(self, status=0, message=None):
         raise ParserExit(status=status, message=message)
+
+    def parse_args(self, args=None, namespace=None):
+        def finish(msg):
+            if self.session and isinstance(self.session, CommandSession):
+                self.session.finish(msg)
+
+        if not args:
+            finish(self.usage)
+        else:
+            try:
+                return super().parse_args(args=args, namespace=namespace)
+            except ParserExit as e:
+                if e.status == 0:
+                    # --help
+                    finish(self.usage)
+                else:
+                    finish('参数不足或不正确，请使用 --help 参数查询使用帮助')