فهرست منبع

--sizes arg; refactor

dennisro 5 سال پیش
والد
کامیت
f2bed58a35
1فایلهای تغییر یافته به همراه21 افزوده شده و 19 حذف شده
  1. 21 19
      2.py

+ 21 - 19
2.py

@@ -15,6 +15,9 @@ optional.add_argument('--threshold', type=str,
 optional.add_argument('--exclude', type=str, nargs='?', action='append',
                     help='Only output files greater than \'size\', e.g. 100M')
 
+optional.add_argument('--sizes', type=str, nargs='+',
+                    help='Only output files greater than \'size\', e.g. 100M')
+
 required.add_argument('--dir', type=str, nargs='?', required=True, action='append',
                     help='Directory to scan. Can be issued multiple times.')
 
@@ -22,12 +25,15 @@ parser._action_groups.append(optional)
 
 args = parser.parse_args()
 
-sizes = ['10M', '50M', '100M', '1G', '5G', 'gt5GB']
+if args.sizes:
+    sizes = args.sizes
+    sizes.append('larger than '+sizes[-1])
+else:
+    sizes = ['10M', '50M', '100M', '1G', '5G']
+    sizes.append('larger than '+sizes[-1])
 
 if args.exclude:
-    print('hi')
     exclude=args.exclude
-    print(exclude)
 
 if args.threshold:
     threshold = humanfriendly.parse_size(args.threshold)
@@ -69,12 +75,11 @@ def joinDicts(dict1, dict2):
 
 def hashfile(path, blocksize = 65536):
     file_size = os.path.getsize(path)
-    # Only hash files larger than threshold
+    # Only hash files larger than threshold (if set)
     if threshold == 0 or (threshold > 0 and file_size > threshold):
         try:
-            print('Hashing '+path)
             afile = open(path, 'rb')
-            hasher = hashlib.sha256()
+            hasher = hashlib.md5()
             buf = afile.read(blocksize)
             while len(buf) > 0:
                 hasher.update(buf)
@@ -100,18 +105,15 @@ def printResults(dict1):
                 final[threshold].append(result)
 
         else:
-            #0=10MB 1=50MB 2=100MB 3=1GB 4=5GB
-            if file_size >= humanfriendly.parse_size(sizes[0]) and file_size < humanfriendly.parse_size(sizes[1]):
-                final[sizes[1]].append(result)
-            elif file_size >= humanfriendly.parse_size(sizes[1]) and file_size < humanfriendly.parse_size(sizes[2]):
-                final[sizes[2]].append(result)
-            elif file_size >= humanfriendly.parse_size(sizes[2]) and file_size < humanfriendly.parse_size(sizes[3]):
-                final[sizes[3]].append(result)
-            elif file_size >= humanfriendly.parse_size(sizes[3]) and file_size < humanfriendly.parse_size(sizes[4]):
-                final[sizes[4]].append(result)
-            elif file_size >= humanfriendly.parse_size(sizes[4]):
-                final[sizes[5]].append(result)
-            else:
+            count = 0
+            while count+1 < len(sizes):
+                try:
+                    if file_size >= humanfriendly.parse_size(sizes[count]) and file_size < humanfriendly.parse_size(sizes[count+1]):
+                        final[sizes[count+1]].append(result)
+                except:
+                    final[sizes[-1]].append(result)
+                count += 1
+            if file_size < humanfriendly.parse_size(sizes[0]):
                 final[sizes[0]].append(result)
             final[threshold]=[False]
     if len(results) > 0 and len(final[threshold]) > 0:
@@ -131,7 +133,7 @@ def printResults(dict1):
             for size in sizes:
                 new.append(size)
                 if len(final[size]) > 0:
-                    if size == 'gt5GB':
+                    if size == 'larger than ' + sizes[-2]:
                         print("\n\033[1;34m\u25b6 >= %s\033[0m" % (new[-2]))
                     else:
                         print("\n\033[1;34m\u25b6 %s to %s\033[0m" % (new[-2],size))