data-collection/data-manager.py at master · abstract-machine-learning/data-collection · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
import os
import sys
import zipfile

# Compresses a single file
def compress(file_path, keep_original = False):
  base_name = os.path.basename(file_path)
  print("Compressing " + file_path + "...")
  with zipfile.ZipFile(file_path + '.zip', 'w', zipfile.ZIP_DEFLATED) as zip_ref:
    zip_ref.write(file_path, base_name)
  if not keep_original:
    os.remove(file_path)


# Recursively compresses every file inside a directory
def compress_all(directory = './domains'):
  for f in os.listdir(directory):
    file_path = directory + '/' + f
    if os.path.isdir(file_path):
      compress_all(file_path)
    elif os.path.isfile(file_path) and os.path.splitext(file_path)[1] == '.csv':
      compress(file_path)


#Decompresses a single zip file
def decompress(file_path):
    base_name = os.path.splitext(os.path.basename(file_path))[0]
    directory = os.path.dirname(file_path)
    print("Decompressing " + file_path + "...")
    with zipfile.ZipFile(file_path, 'r') as zip_ref:
      zip_ref.extractall(directory)


# Recursively decompress every archive in a directory
def decompress_all(directory = './domains'):
  for f in os.listdir(directory):
    file_path = directory + '/' + f
    if os.path.isdir(file_path):
      decompress_all(file_path)
    elif os.path.isfile(file_path) and os.path.splitext(file_path)[1] == '.zip':
      decompress(file_path)


# Recursively removes non-archive files
def clear_all(directory = './domains'):
  for f in os.listdir(directory):
    file_path = directory + '/' + f
    if os.path.isdir(file_path):
      clear_all(file_path)
    elif os.path.isfile(file_path) and os.path.splitext(file_path)[1] != '.zip':
      print("Clearing " + file_path + "...")
      os.remove(file_path)


# Prints usage
def print_usage():
  print("Usage:\n\t" + sys.argv[0] + " compress|decompress|clear <directory>\n")


#######################################################################
# Program entry point

# Input check
if len(sys.argv) < 3:
  print_usage()
  sys.exit()

# Reads mode
mode = sys.argv[1]

# Executes operation
if mode == 'compress':
  compress_all(sys.argv[2])
elif mode == 'decompress':
  decompress_all(sys.argv[2])
elif mode == 'clear':
  clear_all(sys.argv[2])
else:
  print_usage()
  sys.exit()