1 files changed, 10 insertions, 17 deletions
diff --git a/tools/opt-viewer/opt-diff.py b/tools/opt-viewer/opt-diff.py
index 9e921f8488d36..f39432c8bc9dd 100755
--- a/tools/opt-viewer/opt-diff.py
+++ b/tools/opt-viewer/opt-diff.py
@@ -20,24 +20,17 @@ import optrecord
 import argparse
 from collections import defaultdict
 from multiprocessing import cpu_count, Pool
-import os, os.path
-import fnmatch
-
-def find_files(dir_or_file):
-    if os.path.isfile(dir_or_file):
-        return [dir_or_file]
-
-    all = []
-    for dir, subdirs, files in os.walk(dir_or_file):
-        for file in files:
-            if fnmatch.fnmatch(file, "*.opt.yaml"):
-                all.append( os.path.join(dir, file))
-    return all
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description=desc)
-    parser.add_argument('yaml_dir_or_file_1')
-    parser.add_argument('yaml_dir_or_file_2')
+    parser.add_argument(
+        'yaml_dir_or_file_1',
+        help='An optimization record file or a directory searched for optimization '
+             'record files that are used as the old version for the comparison')
+    parser.add_argument(
+        'yaml_dir_or_file_2',
+        help='An optimization record file or a directory searched for optimization '
+             'record files that are used as the new version for the comparison')
     parser.add_argument(
         '--jobs',
         '-j',
@@ -53,8 +46,8 @@ if __name__ == '__main__':
     parser.add_argument('--output', '-o', default='diff.opt.yaml')
     args = parser.parse_args()
 
-    files1 = find_files(args.yaml_dir_or_file_1)
-    files2 = find_files(args.yaml_dir_or_file_2)
+    files1 = optrecord.find_opt_files([args.yaml_dir_or_file_1])
+    files2 = optrecord.find_opt_files([args.yaml_dir_or_file_2])
 
     print_progress = not args.no_progress_indicator
     all_remarks1, _, _ = optrecord.gather_results(files1, args.jobs, print_progress)