Add possibility to write data to file

script/gen_requirements_all.py

@@ -8,6 +8,7 @@ import importlib
 import os
 import pkgutil
 import re
+import argparse
 
 COMMENT_REQUIREMENTS = [
     'RPi.GPIO',
@@ -16,6 +17,7 @@ COMMENT_REQUIREMENTS = [
 
 
 def explore_module(package, explore_children):
+    """ Explore the modules. """
     module = importlib.import_module(package)
 
     found = []
@@ -33,10 +35,10 @@ def explore_module(package, explore_children):
 
 
 def core_requirements():
+    """ Gather core requirements out of setup.py. """
     with open('setup.py') as inp:
         reqs_raw = re.search(
             r'REQUIRES = \[(.*?)\]', inp.read(), re.S).group(1)
-
     return re.findall(r"'(.*?)'", reqs_raw)
 
 
@@ -45,20 +47,23 @@ def comment_requirement(req):
     return any(ign in req for ign in COMMENT_REQUIREMENTS)
 
 
-def main():
+def gather_modules():
-    if not os.path.isfile('requirements_all.txt'):
+    """ Collect the information and construct the output. """
-        print('Run this from HA root dir')
-        return
-
     reqs = OrderedDict()
 
     errors = []
+    output = []
+
     for package in sorted(explore_module('homeassistant.components', True)):
         try:
             module = importlib.import_module(package)
         except ImportError:
             errors.append(package)
             continue
+        # For catching the error by RPi.GPIO
+        # RuntimeError: This module can only be run on a Raspberry Pi!
+        except RuntimeError:
+            continue
 
         if not getattr(module, 'REQUIREMENTS', None):
             continue
@@ -71,20 +76,55 @@ def main():
         print('\n'.join(errors))
         return
 
-    print('# Home Assistant core')
+    output.append('# Home Assistant core')
-    print('\n'.join(core_requirements()))
+    output.append('\n')
-    print()
+    output.append('\n'.join(core_requirements()))
-
+    output.append('\n')
     for pkg, requirements in reqs.items():
         for req in sorted(requirements,
                           key=lambda name: (len(name.split('.')), name)):
-            print('#', req)
+            output.append('\n# {}'.format(req))
 
         if comment_requirement(pkg):
-            print('#', pkg)
+            output.append('\n# {}\n'.format(pkg))
         else:
-            print(pkg)
+            output.append('\n{}\n'.format(pkg))
-        print()
+
+    return ''.join(output)
+
+
+def write_file(data):
+    """ Writes the modules to the requirements_all.txt. """
+    with open('requirements_all.txt', 'w+') as req_file:
+        req_file.write(data)
+
+
+def display(data):
+    """ Prints the output to command line. """
+    print(data)
+
+
+def argparsing():
+    """ Parsing the command line arguments. """
+    parser = argparse.ArgumentParser(
+        description='Generate a requirements_all.txt')
+    parser.add_argument('file', nargs='?',
+                        help='create new requirements_all.txt file')
+    return parser.parse_args()
+
+
+def main():
+    """ Main """
+    if not os.path.isfile('requirements_all.txt'):
+        print('Run this from HA root dir')
+        return
+    args = argparsing()
+    data = gather_modules()
+
+    if args.file:
+        write_file(data)
+    else:
+        display(data)
 
 if __name__ == '__main__':
     main()