Add type hints

1 files changed, 24 insertions, 7 deletions

diff --git a/bin/hostinfo b/bin/hostinfo
index a97a8d8..b85ccc8 100755
--- a/bin/hostinfo
+++ b/bin/hostinfo
@@ -8,6 +8,8 @@ import os
 import pkg_resources
 from dns import resolver, reversename
 
+from typing import Optional
+
 OWN_DIRECTORY = os.path.dirname(os.path.abspath(os.path.realpath(__file__)))
 LIB = os.path.join(OWN_DIRECTORY, '..')
 if os.path.exists(os.path.join(LIB, 'hostinfo')):
@@ -16,10 +18,12 @@ if os.path.exists(os.path.join(LIB, 'hostinfo')):
 from hostinfo import printer
 from hostinfo import utils
 
+
 def _get_data(path):
     stream = file(path, 'r')
     return yaml.load(stream)
 
+
 def _match_key(data, keys):
     if data is None:
         return None
@@ -53,6 +57,7 @@ def _match_key(data, keys):
 
     return None
 
+
 def _match(host, search_key, search_value, negate):
     search_keys = search_key.split('.')
 
@@ -71,6 +76,7 @@ def _match(host, search_key, search_value, negate):
                 return (search_key, result)
         return (search_key, None)
 
+
 def _parse_search(search):
     if search[0] != '?':
         sys.stderr.write("Invalid search string.")
@@ -90,10 +96,13 @@ def _parse_search(search):
 
     return (search, value, negate)
 
-def print_search(basepath, flags, search, filter_key=None):
+
+def print_search(basepath: str, flags: argparse.Namespace,
+                 search: str, filter_key: Optional[str] = None):
     def _get_label(host):
         if flags.short:
-            return host.replace('.spline.inf.fu-berlin.de','')
+            return host.replace('.spline.inf.fu-berlin.de',
+                                '')
         return host
 
     metadata = os.path.join(basepath, 'metadata', 'hosts')
@@ -124,11 +133,13 @@ def print_search(basepath, flags, search, filter_key=None):
                     p.info(key, label=_get_label(host),
                            maxlength=max(length), force=True)
 
+
 def print_info(path, flags, key=None):
     data = _get_data(path)
     p = printer.Printer(data, flags)
     p.info(key)
 
+
 def print_keys(path):
     def _print_keys(data, prefix = ''):
         if isinstance(data, str):
@@ -152,10 +163,11 @@ def print_keys(path):
     data = _get_data(path)
     _print_keys(data)
 
-def print_hosts(path, short):
+
+def print_hosts(path: str, short: bool):
     metadata = os.path.join(path, 'metadata', 'hosts')
     if os.path.exists(metadata):
-        hosts = yaml.load(file(metadata, 'r'))
+        hosts = yaml.safe_load(open(metadata, 'r'))
         if 'hosts' in hosts:
             for host in hosts['hosts']:
                 if short:
@@ -167,7 +179,8 @@ def print_hosts(path, short):
     sys.stderr.write("'%s' not found!\n" % metadata)
     return False
 
-def find_host(basepath, host):
+
+def find_host(basepath: str, host: str):
     path = os.path.join(basepath, host)
     if os.path.exists(path):
         return path
@@ -191,6 +204,7 @@ def find_host(basepath, host):
 
     return None
 
+
 def print_version_and_exit():
     ver = None
     try:
@@ -205,7 +219,8 @@ def print_version_and_exit():
     print(("hostinfo-tools %s" % ver))
     sys.exit(0)
 
-def main():
+
+def main() -> None:
     basepath = '/usr/local/share/hostinfo'
     if 'HOSTINFO_PATH' in os.environ and os.environ['HOSTINFO_PATH'] != '':
         basepath = os.environ['HOSTINFO_PATH']
@@ -256,7 +271,8 @@ def main():
 
     if args.name.startswith('?'):
         # search
-        print_search(basepath, search=args.name, filter_key=args.filter, flags=args)
+        print_search(basepath, search=args.name, filter_key=args.filter,
+                     flags=args)
     else:
         # info
         path = find_host(basepath, args.name)
@@ -273,5 +289,6 @@ def main():
 
     sys.exit(0)
 
+
 if __name__ == '__main__':
     main()