--- a/subsearch.py	Tue Nov 10 20:34:01 2020 +0000
+++ b/subsearch.py	Thu Aug 15 11:00:20 2024 +0000
@@ -13,21 +13,34 @@
 import tempfile
 
 from openbabel import openbabel, pybel
+
 openbabel.obErrorLog.StopLogging()
 
 
 def parse_command_line():
     parser = argparse.ArgumentParser()
-    parser.add_argument('-i', '--infile', required=True, help='Molecule file.')
-    parser.add_argument('--iformat', help='Input format.')
-    parser.add_argument('--fastsearch-index', dest="fastsearch_index", required=True,
-                        help='Path to the openbabel fastsearch index.')
-    parser.add_argument('-o', '--outfile', required=True, help='Path to the output file.')
-    parser.add_argument('--oformat', default='smi', help='Output file format')
-    parser.add_argument("--max-candidates", dest="max_candidates", type=int, default=4000,
-                        help="The maximum number of candidates.")
-    parser.add_argument('-p', '--processors', type=int,
-                        default=multiprocessing.cpu_count())
+    parser.add_argument("-i", "--infile", required=True, help="Molecule file.")
+    parser.add_argument("--iformat", help="Input format.")
+    parser.add_argument(
+        "--fastsearch-index",
+        dest="fastsearch_index",
+        required=True,
+        help="Path to the openbabel fastsearch index.",
+    )
+    parser.add_argument(
+        "-o", "--outfile", required=True, help="Path to the output file."
+    )
+    parser.add_argument("--oformat", default="smi", help="Output file format")
+    parser.add_argument(
+        "--max-candidates",
+        dest="max_candidates",
+        type=int,
+        default=4000,
+        help="The maximum number of candidates.",
+    )
+    parser.add_argument(
+        "-p", "--processors", type=int, default=multiprocessing.cpu_count()
+    )
     return parser.parse_args()
 
 
@@ -40,20 +53,28 @@
 
 def mp_helper(query, args):
     """
-        Helper function for multiprocessing.
-        That function is a wrapper around the following command:
-        obabel file.fs -s"smarts" -Ooutfile.smi -al 999999999
+    Helper function for multiprocessing.
+    That function is a wrapper around the following command:
+    obabel file.fs -s"smarts" -Ooutfile.smi -al 999999999
     """
 
-    if args.oformat == 'names':
-        opts = '-osmi -xt'
+    if args.oformat == "names":
+        opts = "-osmi -xt"
     else:
-        opts = '-o%s' % args.oformat
+        opts = "-o%s" % args.oformat
 
     tmp = tempfile.NamedTemporaryFile(delete=False)
-    cmd = 'obabel -ifs %s -O %s %s -s%s -al %s' % (args.fastsearch_index, tmp.name, opts, query, args.max_candidates)
+    cmd = "obabel -ifs %s -O %s %s -s%s -al %s" % (
+        args.fastsearch_index,
+        tmp.name,
+        opts,
+        query,
+        args.max_candidates,
+    )
 
-    child = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    child = subprocess.Popen(
+        cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    )
 
     stdout, stderr = child.communicate()
     return_code = child.returncode
@@ -73,14 +94,14 @@
     """
     Wrapper to retrieve a striped SMILES or SMARTS string from different input formats.
     """
-    if args.iformat in ['smi', 'text', 'tabular']:
+    if args.iformat in ["smi", "text", "tabular"]:
         with open(args.infile) as text_file:
             for line in text_file:
-                yield line.split('\t')[0].strip()
+                yield line.split("\t")[0].strip()
     else:
         # inchi or sdf files
         for mol in pybel.readfile(args.iformat, args.infile):
-            yield mol.write('smiles').split('\t')[0]
+            yield mol.write("smiles").split("\t")[0]
 
 
 def substructure_search(args):
@@ -91,18 +112,18 @@
     pool.close()
     pool.join()
 
-    if args.oformat == 'names':
-        out_handle = open(args.outfile, 'w')
+    if args.oformat == "names":
+        out_handle = open(args.outfile, "w")
         for result_file, query in results:
             with open(result_file) as res_handle:
                 for line in res_handle:
-                    out_handle.write('%s\t%s\n' % (line.strip(), query))
+                    out_handle.write("%s\t%s\n" % (line.strip(), query))
             os.remove(result_file)
         out_handle.close()
     else:
-        out_handle = open(args.outfile, 'wb')
+        out_handle = open(args.outfile, "wb")
         for result_file, query in results:
-            res_handle = open(result_file, 'rb')
+            res_handle = open(result_file, "rb")
             shutil.copyfileobj(res_handle, out_handle)
             res_handle.close()
             os.remove(result_file)
@@ -111,7 +132,7 @@
 
 def __main__():
     """
-        Multiprocessing Open Babel Substructure Search.
+    Multiprocessing Open Babel Substructure Search.
     """
     args = parse_command_line()
     substructure_search(args)