Merge pull request #109 from gOATiful/master

Thanks @gOATiful !

Author: urialon

JavaExtractor/extract.py

@@ -12,32 +12,27 @@
 from subprocess import Popen, PIPE, STDOUT, call
 
 
-
 def get_immediate_subdirectories(a_dir):
     return [(os.path.join(a_dir, name)) for name in os.listdir(a_dir)
             if os.path.isdir(os.path.join(a_dir, name))]
 
 
-TMP_DIR = ""
-
-def ParallelExtractDir(args, dir):
-    ExtractFeaturesForDir(args, dir, "")
+def ParallelExtractDir(args, tmpdir, dir_):
+    ExtractFeaturesForDir(args,tmpdir, dir_, "")
 
 
-def ExtractFeaturesForDir(args, dir, prefix):
+def ExtractFeaturesForDir(args, tmpdir, dir_, prefix):
     command = ['java', '-cp', args.jar, 'JavaExtractor.App',
                '--max_path_length', str(args.max_path_length), '--max_path_width', str(args.max_path_width),
-               '--dir', dir, '--num_threads', str(args.num_threads)]
-
+               '--dir', dir_, '--num_threads', str(args.num_threads)]
     # print command
     # os.system(command)
     kill = lambda process: process.kill()
-    outputFileName = TMP_DIR + prefix + dir.split('/')[-1]
+    outputFileName = tmpdir + prefix + dir_.split('/')[-1]
     failed = False
     with open(outputFileName, 'a') as outputFile:
-        sleeper = subprocess.Popen(command, stdout=outputFile, stderr=subprocess.PIPE)
+        sleeper = subprocess.Popen(command, stdout=outputFile, stderr=subprocess.PIPE,)
         timer = Timer(600000, kill, [sleeper])
-
         try:
             timer.start()
             stdout, stderr = sleeper.communicate()
@@ -48,32 +43,31 @@ def ExtractFeaturesForDir(args, dir, prefix):
             if len(stderr) > 0:
                 print(sys.stderr, stderr, file=sys.stdout)
         else:
-            print(sys.stderr, 'dir: ' + str(dir) + ' was not completed in time', file=sys.stdout)
+            print(sys.stderr, 'dir: ' + str(dir_) + ' was not completed in time', file=sys.stdout, flush=True)
             failed = True
-            subdirs = get_immediate_subdirectories(dir)
+            subdirs = get_immediate_subdirectories(dir_)
             for subdir in subdirs:
-                ExtractFeaturesForDir(args, subdir, prefix + dir.split('/')[-1] + '_')
+                ExtractFeaturesForDir(args, subdir, prefix + dir_.split('/')[-1] + '_')
     if failed:
         if os.path.exists(outputFileName):
             os.remove(outputFileName)
 
 
 def ExtractFeaturesForDirsList(args, dirs):
-    global TMP_DIR
-    TMP_DIR = "./tmp/feature_extractor%d/" % (os.getpid())
-    if os.path.exists(TMP_DIR):
-        shutil.rmtree(TMP_DIR, ignore_errors=True)
-    os.makedirs(TMP_DIR)
+    tmp_dir = f"./tmp/feature_extractor{os.getpid()}/"
+    if os.path.exists(tmp_dir):
+        shutil.rmtree(tmp_dir, ignore_errors=True)
+    os.makedirs(tmp_dir)
     try:
         p = multiprocessing.Pool(4)
-        p.starmap(ParallelExtractDir, zip(itertools.repeat(args), dirs))
+        p.starmap(ParallelExtractDir, zip(itertools.repeat(args),itertools.repeat(tmp_dir), dirs))
         #for dir in dirs:
         #    ExtractFeaturesForDir(args, dir, '')
-        output_files = os.listdir(TMP_DIR)
+        output_files = os.listdir(tmp_dir)
         for f in output_files:
-            os.system("cat %s/%s" % (TMP_DIR, f))
+            os.system("cat %s/%s" % (tmp_dir, f))
     finally:
-        shutil.rmtree(TMP_DIR, ignore_errors=True)
+        shutil.rmtree(tmp_dir, ignore_errors=True)
 
 
 if __name__ == '__main__':