pgtoolkit/toolbase.py
changeset 100 d6088dba8fea
parent 84 3b5dd9efba35
child 101 2a2d0d5df03b
--- a/pgtoolkit/toolbase.py	Tue May 06 18:34:38 2014 +0200
+++ b/pgtoolkit/toolbase.py	Tue May 06 18:37:41 2014 +0200
@@ -1,12 +1,13 @@
-import argparse
-import logging
-import re
+from pgtoolkit import pgmanager, pgbrowser
 
 from pycolib.configparser import ConfigParser
 from pycolib.coloredformatter import ColoredFormatter
 from pycolib.ansicolor import highlight
 
-from pgtoolkit import pgmanager, pgbrowser
+import argparse
+import logging
+import re
+import textwrap
 
 
 class ConnectionInfoNotFound(Exception):
@@ -17,9 +18,17 @@
     pass
 
 
+class ToolDescriptionFormatter(argparse.HelpFormatter):
+    """Help message formatter which retains any formatting in descriptions."""
+
+    def _fill_text(self, text, width, indent):
+        return textwrap.dedent(text)
+
+
 class ToolBase:
     def __init__(self, name, desc, **kwargs):
-        self.parser = argparse.ArgumentParser(description=desc)
+        self.parser = argparse.ArgumentParser(prog=name, description=desc,
+            formatter_class=ToolDescriptionFormatter)
         self.parser.add_argument('-d', dest='debug', action='store_true',
             help='Debug mode - print database queries.')
 
@@ -31,9 +40,9 @@
         self.pgm = pgmanager.get_instance()
         self.target_isolation_level = None
 
-    def init(self):
+    def init(self, args=None):
         self.config.load('pgtoolkit.conf')
-        self.args = self.parser.parse_args()
+        self.args = self.parser.parse_args(args)
         self.init_logging()
 
     def init_logging(self):
@@ -110,8 +119,8 @@
         ToolBase.__init__(self, name, desc, **kwargs)
         self.parser.add_argument('target', metavar='target', type=str, help='Target database')
 
-    def init(self):
-        ToolBase.init(self)
+    def init(self, args=None):
+        ToolBase.init(self, args)
         self.prepare_conns(target=self.args.target)
 
 
@@ -123,8 +132,8 @@
         if 'allow_reverse' in kwargs and kwargs['allow_reverse']:
             self.parser.add_argument('-r', '--reverse', action='store_true', help='Reverse operation. Swap source and destination.')
 
-    def init(self):
-        ToolBase.init(self)
+    def init(self, args=None):
+        ToolBase.init(self, args)
         if self.is_reversed():
             self.args.src, self.args.dst = self.args.dst, self.args.src
         self.prepare_conns(src=self.args.src, dst=self.args.dst)
@@ -146,8 +155,8 @@
             dest='dstschema', type=str, default='', help='Destination schema name (default=source_schema).')
         self.parser.add_argument('--regex', action='store_true', help="Use RE in schema or table name.")
 
-    def init(self):
-        SrcDstTool.init(self)
+    def init(self, args=None):
+        SrcDstTool.init(self, args)
 
         self.schema1 = self.args.srcschema
         self.table1 = self.args.srctable