Fix argparse usage and error reporting.

bdbcontrib/contrib.py

@@ -14,13 +14,24 @@
 import matplotlib
 matplotlib.rcParams.update({'figure.autolayout': True})
 import matplotlib.pyplot as plt
-import seaborn as sns
+# import seaborn as sns
 
 
 ROOTDIR = os.path.dirname(os.path.abspath(__file__))
 READTOHTML_CSS = os.path.join(ROOTDIR, 'readtohtml.css')
 
 
+class ArgparseError(Exception):
+    def __init__(self, status, message):
+        self.status = status
+        self.message = message
+
+
+class ArgumentParser(argparse.ArgumentParser):
+    def exit(self, status=0, message=None):
+        raise ArgparseError(status, message)
+
+
 @bayesdb_shell_cmd('register_bql_math_functions')
 def register_bql_math(self, args):
     '''Adds basic math functions to BQL'''
@@ -94,15 +105,19 @@ def zmatrix(self, argin):
     Example:
     bayeslite> .heatmap ESTIMATE PAIRWISE DEPENDENCE PROBABILITY FROM mytable
     '''
-    parser = argparse.ArgumentParser()
+    parser = ArgumentParser(prog='.heatmap')
     parser.add_argument('bql', type=str, nargs='+', help='PAIRWISE BQL query')
     parser.add_argument('-f', '--filename', type=str, default=None,
                         help='output filename')
     parser.add_argument('--vmin', type=float, default=None)
     parser.add_argument('--vmax', type=float, default=None)
     parser.add_argument('--last-sort', action='store_true')
 
-    args = parser.parse_args(shlex.split(argin))
+    try:
+        args = parser.parse_args(shlex.split(argin))
+    except ArgparseError as e:
+        self.stdout.write('%s' % (e.message,))
+        return
 
     bql = " ".join(args.bql)
 
@@ -166,7 +181,7 @@ def pairplot(self, argin):
         --filename myfile.png
     bayeslite> .show ESTIMATE foo, baz FROM mytable_cc -g mytable_cc
     '''
-    parser = argparse.ArgumentParser()
+    parser = ArgumentParser(prog='.show')
     parser.add_argument('bql', type=str, nargs='+', help='BQL query')
     parser.add_argument('-f', '--filename', type=str, default=None,
                         help='output filename')
@@ -180,7 +195,11 @@ def pairplot(self, argin):
                         help='Plot missing values in scatterplot.')
     parser.add_argument('--colorby', type=str, default=None,
                         help='Name of column to use as a dummy variable.')
-    args = parser.parse_args(shlex.split(argin))
+    try:
+        args = parser.parse_args(shlex.split(argin))
+    except ArgparseError as e:
+        self.stdout.write('%s' % (e.message,))
+        return
 
     bql = " ".join(args.bql)
 
@@ -252,15 +271,19 @@ def histogram(self, argin):
     bayeslite> .histogram SELECT height, sex FROM humans; --normed --bin 31
     '''
 
-    parser = argparse.ArgumentParser()
+    parser = ArgumentParser(prog='.histogram')
     parser.add_argument('bql', type=str, nargs='+', help='BQL query')
     parser.add_argument('-f', '--filename', type=str, default=None,
                         help='output filename')
     parser.add_argument('-b', '--bins', type=int, default=15,
                         help='number of bins')
     parser.add_argument('--normed', action='store_true',
                         help='Normalize histograms?')
-    args = parser.parse_args(shlex.split(argin))
+    try:
+        args = parser.parse_args(shlex.split(argin))
+    except ArgparseError as e:
+        self.stdout.write('%s' % (e.message,))
+        return
 
     bql = " ".join(args.bql)
 
@@ -277,11 +300,15 @@ def histogram(self, argin):
 @bayesdb_shell_cmd('bar')
 def barplot(self, argin):
     '''FIXME'''
-    parser = argparse.ArgumentParser()
+    parser = ArgumentParser('.bar')
     parser.add_argument('bql', type=str, nargs='+', help='BQL query')
     parser.add_argument('-f', '--filename', type=str, default=None,
                         help='output filename')
-    args = parser.parse_args(shlex.split(argin))
+    try:
+        args = parser.parse_args(shlex.split(argin))
+    except ArgparseError as e:
+        self.stdout.write('%s' % (e.message,))
+        return
     bql = " ".join(args.bql)
 
     df = do_query(self._bdb, bql).as_df()