add root argument

Author: BruceZhang1993

index_generator/__main__.py

@@ -23,6 +23,7 @@ def main():
                         help='Default output filename.')
     parser.add_argument('--print', '-P', action='store_true', default=False, help='Whether to print to stdout.')
     parser.add_argument('--depth', '-d', type=int, default=0, help='Set cutoff depth.')
+    parser.add_argument('--root', '-r', type=str, default='/', help='Set base root dir.')
     parser.add_argument('path', type=str, default='', help='Path', nargs='?')
     arguments = parser.parse_args()
     app(arguments)
@@ -38,19 +39,19 @@ def app(args):
         print('See: index-generator --help')
         sys.exit(0)
     if args.no_recursive:
-        generate_once(args.template, args.path, os.listdir(args.path), args.name, args.print)
+        generate_once(args.template, args.path, os.listdir(args.path), args.name, args.print, base=args.root)
     else:
-        generate_recursively(args.template, args.path, args.name, args.print, args.depth)
+        generate_recursively(args.template, args.path, args.name, args.print, args.depth, base=args.root)
 
 
-def generate_once(template_dir, root, files, name, if_print):
+def generate_once(template_dir, root, files, name, if_print, base='/'):
     environment = jinja2.Environment(
         loader=jinja2.PackageLoader('index_generator', template_dir),
         autoescape=jinja2.select_autoescape(['html', 'htm'])
     )
     template = environment.get_template(name)
 
-    entries = list(map(lambda f1: Entry(f1, root), files))
+    entries = list(map(lambda f1: Entry(f1, root, base=base), files))
     # entries.sort(key=lambda x: x.isDir, reverse=True)
 
     filelist = []
@@ -66,7 +67,7 @@ def generate_once(template_dir, root, files, name, if_print):
             'isDir':    entry.isDir
         })
     html = template.render(ig={
-        'root': '/'+root.lstrip('.*/'),
+        'root': base + root.lstrip('.*/'),
         'files': filelist,
         'generator': {
             'name':    APP_NAME,
@@ -78,11 +79,11 @@ def generate_once(template_dir, root, files, name, if_print):
     if if_print:
         print(html)
     else:
-        with open(root+os.path.sep+name, 'w') as f:
+        with open(root + os.path.sep + name, 'w') as f:
             print(html, file=f)
 
 
-def generate_recursively(template_dir, path, name, if_print, max_depth=0):
+def generate_recursively(template_dir, path, name, if_print, max_depth=0, base='/'):
     for root, dirs, files in os.walk(path):
         if max_depth != 0 and root.count(os.sep) >= max_depth:
             dirs.clear()
@@ -98,7 +99,7 @@ def generate_recursively(template_dir, path, name, if_print, max_depth=0):
             print('files: {}'.format(files))
             print('-----------------------------------------')
 
-        generate_once(template_dir, root, dirs+files, name, if_print)
+        generate_once(template_dir, root, dirs+files, name, if_print, base=base)
 
 
 if __name__ == '__main__':

index_generator/models/entries.py

@@ -3,9 +3,9 @@
 
 
 class Entry(object):
-    def __init__(self, file, root):
+    def __init__(self, file, root, base='/'):
         path = root + os.path.sep + file
-        self.path = '/' + path.lstrip('.*/')
+        self.path = base + path.lstrip('.*/')
         self.name = os.path.basename(path)
         self.mime = mimetypes.guess_type(path)[0]
         self.size = os.path.getsize(path)