taxaforge 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) [2024] [Pandikunta Anand Reddy]
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,93 @@
1
+ Metadata-Version: 2.4
2
+ Name: taxaforge
3
+ Version: 0.2.0
4
+ Summary: CLI tool to build taxonomic classifier databases (Kraken2, ganon2, ...) with a single command
5
+ Maintainer-email: ChillarAnand <anand@avilpage.com>
6
+ Requires-Python: >=3.8
7
+ Description-Content-Type: text/markdown
8
+ Classifier: Development Status :: 4 - Beta
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: License :: OSI Approved :: MIT License
11
+ Classifier: Operating System :: OS Independent
12
+ Classifier: Programming Language :: Python
13
+ Classifier: Typing :: Typed
14
+ Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
15
+ Classifier: Topic :: Utilities
16
+ License-File: LICENSE.txt
17
+ Requires-Dist: ncbi-genome-download
18
+ Requires-Dist: click
19
+ Project-URL: Documentation, https://avilpage.com/taxaforge.html
20
+ Project-URL: Source, https://github.com/AvilPage/TaxaForge
21
+
22
+ Installation
23
+ ============
24
+
25
+ ```bash
26
+ pip install taxaforge
27
+ ```
28
+
29
+ Usage
30
+ =====
31
+
32
+ ```bash
33
+ taxaforge --help
34
+ ```
35
+
36
+ To create standard Kraken2 database
37
+
38
+ ```bash
39
+ taxaforge build --db-type standard
40
+ ```
41
+
42
+ Before creating a standard database, you can try a smaller database like fungi.
43
+
44
+ ```bash
45
+ taxaforge build --db-type fungi
46
+ ```
47
+
48
+ To build a ganon2 database instead, pass `--tool ganon2`
49
+
50
+ ```bash
51
+ taxaforge build --tool ganon2 --db-type fungi
52
+ ```
53
+
54
+ To use locally downloaded files, run the following command
55
+
56
+ ```bash
57
+ taxaforge build --db-name k2_test --genomes-dir /path/to/genomes --taxonomy-dir /path/to/taxonomy
58
+ ```
59
+
60
+ To limit the number of genomes in the database, use the `--limit` option
61
+
62
+ ```bash
63
+ taxaforge build --db-name k2_test_100 --genomes-dir /path/to/genomes --limit 1000
64
+ ```
65
+
66
+ Config
67
+ ======
68
+
69
+ Read/write config, stored in the OS default config location (`~/.config/taxaforge/config.ini` on Linux, `~/Library/Application Support/taxaforge/config.ini` on macOS).
70
+
71
+ ```bash
72
+ taxaforge config set threads 8
73
+ taxaforge config get threads
74
+ taxaforge config
75
+ ```
76
+
77
+
78
+ Why TaxaForge?
79
+ ==============
80
+
81
+ TaxaForge aims to provide a simple and easy to use tool to build wide variety of taxonomic classifier databases with a single command.
82
+
83
+ Why not kraken2-build/ganon directly?
84
+
85
+ kraken2-build and ganon each build databases for their own tool only. TaxaForge wraps the shared download/taxonomy pipeline once and dispatches to the right tool via `--tool`, so adding support for more classifiers is a matter of plugging in a new build step.
86
+
87
+
88
+ Documentation
89
+ =============
90
+
91
+ - [Kraken2 Database Builder](https://avilpage.com/kdb.html)
92
+ - [Mastering Kraken2](https://avilpage.com/tags/kraken2.html)
93
+
@@ -0,0 +1,71 @@
1
+ Installation
2
+ ============
3
+
4
+ ```bash
5
+ pip install taxaforge
6
+ ```
7
+
8
+ Usage
9
+ =====
10
+
11
+ ```bash
12
+ taxaforge --help
13
+ ```
14
+
15
+ To create standard Kraken2 database
16
+
17
+ ```bash
18
+ taxaforge build --db-type standard
19
+ ```
20
+
21
+ Before creating a standard database, you can try a smaller database like fungi.
22
+
23
+ ```bash
24
+ taxaforge build --db-type fungi
25
+ ```
26
+
27
+ To build a ganon2 database instead, pass `--tool ganon2`
28
+
29
+ ```bash
30
+ taxaforge build --tool ganon2 --db-type fungi
31
+ ```
32
+
33
+ To use locally downloaded files, run the following command
34
+
35
+ ```bash
36
+ taxaforge build --db-name k2_test --genomes-dir /path/to/genomes --taxonomy-dir /path/to/taxonomy
37
+ ```
38
+
39
+ To limit the number of genomes in the database, use the `--limit` option
40
+
41
+ ```bash
42
+ taxaforge build --db-name k2_test_100 --genomes-dir /path/to/genomes --limit 1000
43
+ ```
44
+
45
+ Config
46
+ ======
47
+
48
+ Read/write config, stored in the OS default config location (`~/.config/taxaforge/config.ini` on Linux, `~/Library/Application Support/taxaforge/config.ini` on macOS).
49
+
50
+ ```bash
51
+ taxaforge config set threads 8
52
+ taxaforge config get threads
53
+ taxaforge config
54
+ ```
55
+
56
+
57
+ Why TaxaForge?
58
+ ==============
59
+
60
+ TaxaForge aims to provide a simple and easy to use tool to build wide variety of taxonomic classifier databases with a single command.
61
+
62
+ Why not kraken2-build/ganon directly?
63
+
64
+ kraken2-build and ganon each build databases for their own tool only. TaxaForge wraps the shared download/taxonomy pipeline once and dispatches to the right tool via `--tool`, so adding support for more classifiers is a matter of plugging in a new build step.
65
+
66
+
67
+ Documentation
68
+ =============
69
+
70
+ - [Kraken2 Database Builder](https://avilpage.com/kdb.html)
71
+ - [Mastering Kraken2](https://avilpage.com/tags/kraken2.html)
@@ -0,0 +1,94 @@
1
+ [project]
2
+ name = "taxaforge"
3
+ version = "0.2.0"
4
+ description = "CLI tool to build taxonomic classifier databases (Kraken2, ganon2, ...) with a single command"
5
+ readme = "README.md"
6
+ license = {file = "LICENSE.txt"}
7
+ maintainers = [{name = "ChillarAnand", email = "anand@avilpage.com"}]
8
+ classifiers = [
9
+ "Development Status :: 4 - Beta",
10
+ "Intended Audience :: Developers",
11
+ "License :: OSI Approved :: MIT License",
12
+ "Operating System :: OS Independent",
13
+ "Programming Language :: Python",
14
+ "Typing :: Typed",
15
+ "Topic :: Scientific/Engineering :: Bio-Informatics",
16
+ "Topic :: Utilities",
17
+ ]
18
+ requires-python = ">=3.8"
19
+ dependencies = [
20
+ "ncbi-genome-download",
21
+ "click",
22
+ ]
23
+
24
+ [project.urls]
25
+ Documentation = "https://avilpage.com/taxaforge.html"
26
+ Source = "https://github.com/AvilPage/TaxaForge"
27
+
28
+ [build-system]
29
+ requires = ["flit_core<4"]
30
+ build-backend = "flit_core.buildapi"
31
+
32
+ [tool.flit.module]
33
+ name = "taxaforge"
34
+
35
+ [tool.flit.sdist]
36
+ include = [
37
+ "requirements/",
38
+ "tests/",
39
+ ]
40
+ exclude = [
41
+ "docs/_build/",
42
+ ]
43
+
44
+ [tool.pytest.ini_options]
45
+ testpaths = ["tests"]
46
+ filterwarnings = [
47
+ "error",
48
+ ]
49
+
50
+ [tool.coverage.run]
51
+ branch = true
52
+ source = ["click", "tests"]
53
+
54
+ [tool.coverage.paths]
55
+ source = ["src", "*/site-packages"]
56
+
57
+ [tool.mypy]
58
+ python_version = "3.8"
59
+ files = ["src/click", "tests/typing"]
60
+ show_error_codes = true
61
+ pretty = true
62
+ strict = true
63
+
64
+ ignore_missing_imports = true
65
+
66
+ [tool.pyright]
67
+ pythonVersion = "3.8"
68
+ include = ["src/click", "tests/typing"]
69
+ typeCheckingMode = "basic"
70
+
71
+ [tool.ruff]
72
+ extend-exclude = ["examples/"]
73
+ src = ["src"]
74
+ fix = true
75
+ show-fixes = true
76
+ output-format = "full"
77
+
78
+ [tool.ruff.lint]
79
+ select = [
80
+ "B", # flake8-bugbear
81
+ "E", # pycodestyle error
82
+ "F", # pyflakes
83
+ "I", # isort
84
+ "UP", # pyupgrade
85
+ "W", # pycodestyle warning
86
+ ]
87
+ ignore-init-module-imports = true
88
+
89
+ [tool.ruff.lint.isort]
90
+ force-single-line = true
91
+ order-by-type = false
92
+
93
+ [project.scripts]
94
+ taxaforge = "taxaforge:cli"
@@ -0,0 +1,529 @@
1
+ #!/usr/bin/env python3
2
+ import concurrent.futures
3
+ import configparser
4
+ import datetime
5
+ import hashlib
6
+ import logging
7
+ import multiprocessing
8
+ import os
9
+ import shlex
10
+ import shutil
11
+ import subprocess
12
+ import sys
13
+ import urllib.request
14
+ from pathlib import Path
15
+
16
+ import click
17
+ import ncbi_genome_download
18
+ from tqdm import tqdm
19
+
20
+ logger = logging.getLogger(__name__)
21
+ logger.setLevel(logging.INFO)
22
+ logger.addHandler(logging.StreamHandler())
23
+
24
+
25
+ NCBI_SERVER = "https://ftp.ncbi.nlm.nih.gov"
26
+
27
+
28
+ DB_TYPE_CONFIG = {
29
+ 'standard': ("archaea", "bacteria", "viral", "plasmid", "human", "UniVec_Core")
30
+ }
31
+ REQUIRED_BINS = {
32
+ 'kraken2': "kraken2-build",
33
+ 'ganon2': "ganon",
34
+ 'ganon': "ganon",
35
+ }
36
+ hashes = set()
37
+ md5_file = None
38
+
39
+
40
+ def hash_file(filename, buf_size=8192):
41
+ md5 = hashlib.md5()
42
+ with open(filename, "rb") as in_file:
43
+ while True:
44
+ data = in_file.read(buf_size)
45
+ if not data:
46
+ break
47
+ md5.update(data)
48
+ digest = md5.hexdigest()
49
+ return digest
50
+
51
+
52
+ def run_basic_checks(tool, use_k2=False):
53
+ if not shutil.which("ncbi-genome-download"):
54
+ logger.error("ncbi-genome-download not found in PATH. Exiting.")
55
+ sys.exit(1)
56
+
57
+ if tool not in REQUIRED_BINS:
58
+ logger.error(f"Unknown tool: {tool}. Supported tools: {', '.join(REQUIRED_BINS)}")
59
+ sys.exit(1)
60
+
61
+ binary = "k2" if (tool == 'kraken2' and use_k2) else REQUIRED_BINS[tool]
62
+ if not shutil.which(binary):
63
+ logger.error(f"{binary} not found in PATH. Exiting.")
64
+ sys.exit(1)
65
+
66
+
67
+ def create_cache_dir():
68
+ # Unix ~/.cache/taxaforge
69
+ # macOS ~/Library/Caches/taxaforge
70
+ if sys.platform == "darwin":
71
+ cache_dir = Path.home() / "Library" / "Caches" / "taxaforge"
72
+ if sys.platform == "linux":
73
+ cache_dir = Path.home() / ".cache" / "taxaforge"
74
+
75
+ cache_dir.mkdir(parents=True, exist_ok=True)
76
+ return cache_dir
77
+
78
+
79
+ def create_config_dir():
80
+ # Unix ~/.config/taxaforge
81
+ # macOS ~/Library/Application Support/taxaforge
82
+ if sys.platform == "darwin":
83
+ config_dir = Path.home() / "Library" / "Application Support" / "taxaforge"
84
+ if sys.platform == "linux":
85
+ config_dir = Path.home() / ".config" / "taxaforge"
86
+
87
+ config_dir.mkdir(parents=True, exist_ok=True)
88
+ return config_dir
89
+
90
+
91
+ CONFIG_SECTION = "taxaforge"
92
+
93
+
94
+ def get_config_path():
95
+ return create_config_dir() / "config.ini"
96
+
97
+
98
+ def load_config():
99
+ parser = configparser.ConfigParser()
100
+ parser.read(get_config_path())
101
+ if not parser.has_section(CONFIG_SECTION):
102
+ parser.add_section(CONFIG_SECTION)
103
+ return parser
104
+
105
+
106
+ def save_config(parser):
107
+ with open(get_config_path(), "w") as out_file:
108
+ parser.write(out_file)
109
+
110
+
111
+ def download_file(url, position):
112
+ filename = url.rsplit("/", 1)[-1]
113
+ existing = os.path.getsize(filename) if os.path.exists(filename) else 0
114
+
115
+ request = urllib.request.Request(url)
116
+ if existing:
117
+ request.add_header("Range", f"bytes={existing}-")
118
+
119
+ with urllib.request.urlopen(request) as response:
120
+ resumed = response.status == 206
121
+ total = int(response.headers.get("Content-Length", 0)) + (existing if resumed else 0)
122
+
123
+ with open(filename, "ab" if resumed else "wb") as out_file, tqdm(
124
+ total=total, initial=existing if resumed else 0, unit="B", unit_scale=True,
125
+ desc=filename, position=position, leave=True
126
+ ) as bar:
127
+ while True:
128
+ chunk = response.read(1024 * 1024)
129
+ if not chunk:
130
+ break
131
+ out_file.write(chunk)
132
+ bar.update(len(chunk))
133
+
134
+
135
+ def download_files(urls, max_workers=4):
136
+ with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
137
+ futures = [
138
+ executor.submit(download_file, url, index % max_workers)
139
+ for index, url in enumerate(urls)
140
+ ]
141
+ for future in concurrent.futures.as_completed(futures):
142
+ future.result()
143
+
144
+
145
+ def download_taxanomy(cache_dir, skip_maps=None, protein=None):
146
+ taxonomy_path = os.path.join(cache_dir, "taxonomy")
147
+ os.makedirs(taxonomy_path, exist_ok=True)
148
+ os.chdir(taxonomy_path)
149
+
150
+ if not skip_maps:
151
+ if not protein:
152
+ # Define URLs for nucleotide accession to taxon map
153
+ urls = [
154
+ f"{NCBI_SERVER}/pub/taxonomy/accession2taxid/nucl_gb.accession2taxid.gz",
155
+ f"{NCBI_SERVER}/pub/taxonomy/accession2taxid/nucl_wgs.accession2taxid.gz"
156
+ ]
157
+ else:
158
+ # Define URL for protein accession to taxon map
159
+ urls = ["ftp://ftp.ncbi.nlm.nih.gov/pub/taxonomy/accession2taxid/prot.accession2taxid.gz"]
160
+ else:
161
+ logger.info("Skipping maps download")
162
+
163
+ # Download taxonomy tree data
164
+ urls.append(f"{NCBI_SERVER}/pub/taxonomy/taxdump.tar.gz")
165
+
166
+ logger.info(f"Downloading {len(urls)} taxonomy files")
167
+ download_files(urls)
168
+
169
+ logger.info("Extracting taxdump.tar.gz")
170
+ cmd = f"tar -k -xvf taxdump.tar.gz"
171
+ run_cmd(cmd)
172
+
173
+ logger.info("Decompressing taxonomy data")
174
+ cmd = f"find {cache_dir}/taxonomy -name '*.gz' | xargs -n 1 gunzip -k"
175
+ run_cmd(cmd)
176
+
177
+ logger.info("Finished downloading taxonomy data")
178
+
179
+
180
+ def run_cmd(cmd, return_output=False, no_output=False):
181
+ if not no_output:
182
+ logger.info(f"Running command: {cmd}")
183
+
184
+ if return_output:
185
+ return subprocess.check_output(cmd, shell=True).decode("utf-8").strip().split("\n")
186
+
187
+ try:
188
+ if no_output:
189
+ subprocess.run(cmd, shell=True, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
190
+ else:
191
+ subprocess.run(cmd, shell=True, check=True)
192
+ except subprocess.CalledProcessError:
193
+ pass
194
+
195
+
196
+ def download_genomes(cache_dir, cwd, db_type, db_name, threads, force=False):
197
+ organisms = DB_TYPE_CONFIG.get(db_type, [db_type])
198
+ if force:
199
+ shutil.rmtree(cwd / db_name, ignore_errors=True)
200
+
201
+ os.makedirs(cwd / db_name, exist_ok=True)
202
+
203
+ for organism in organisms:
204
+ logger.info(f"Downloading genomes for {organism}")
205
+ os.chdir(cache_dir)
206
+ ncbi_genome_download.download(
207
+ section='refseq', groups=organism, file_formats='fasta',
208
+ progress_bar=True, parallel=threads,
209
+ assembly_levels=['complete'],
210
+ output=cache_dir
211
+ )
212
+
213
+ cmd = f"find {cache_dir}/refseq/{organism} -name '*.gz' | xargs -n 1 -P {threads} gunzip -k"
214
+ run_cmd(cmd)
215
+ logger.info(f"Finished downloading {organism} genomes")
216
+
217
+ os.chdir(cwd)
218
+ logger.info("Finished downloading all genomes")
219
+
220
+
221
+ def build_db(
222
+ cache_dir, cwd, db_type, db_name, threads, kmer_len, min_len,
223
+ fast_build, rebuild, load_factor, use_k2
224
+ ):
225
+ run_cmd(f"cd {cwd}")
226
+
227
+ if not os.path.exists(f"{db_name}/taxonomy"):
228
+ cmd = f"ln -s {cache_dir}/taxonomy {db_name}/"
229
+ run_cmd(cmd)
230
+
231
+ if rebuild:
232
+ cmd = f"rm -rf {db_name}/*.k2d"
233
+ run_cmd(cmd)
234
+
235
+ # TODO: Fix issue with macos threads
236
+ if sys.platform == "darwin":
237
+ threads = 1
238
+
239
+ if use_k2:
240
+ cmd = f"k2 build"
241
+ else:
242
+ cmd = f"kraken2-build --build"
243
+
244
+ cmd += f" --db {db_name} --threads {threads} --kmer-len {kmer_len} --minimizer-len {min_len} --load-factor {load_factor}"
245
+ if fast_build:
246
+ cmd += " --fast-build"
247
+
248
+ run_cmd(cmd)
249
+
250
+ cmd = f"du -sh {db_name}/*.k2d"
251
+ run_cmd(cmd)
252
+
253
+
254
+ def build_ganon2(cache_dir, cwd, genomes_dir, db_type, db_name, threads, kmer_len, min_len, level, rebuild):
255
+ os.chdir(cwd)
256
+
257
+ if genomes_dir:
258
+ input_dirs = [str(genomes_dir)]
259
+ else:
260
+ organisms = DB_TYPE_CONFIG.get(db_type, [db_type])
261
+ input_dirs = [f"{cache_dir}/refseq/{organism}" for organism in organisms]
262
+
263
+ if rebuild:
264
+ cmd = f"rm -f {db_name}.*"
265
+ run_cmd(cmd)
266
+
267
+ cmd = (
268
+ f"ganon build-custom --input {' '.join(input_dirs)} --input-recursive "
269
+ f"--taxonomy-files {cache_dir}/taxonomy/nodes.dmp {cache_dir}/taxonomy/names.dmp "
270
+ f"--db-prefix {db_name} --threads {threads} --kmer-size {kmer_len} "
271
+ f"--window-size {min_len} --level {level}"
272
+ )
273
+ run_cmd(cmd)
274
+
275
+ cmd = f"du -sh {db_name}.*"
276
+ run_cmd(cmd)
277
+
278
+
279
+ def get_files(genomes_dir, cache_dir, db_type, db_name, threads):
280
+ if genomes_dir:
281
+ logger.info(f"Adding {genomes_dir} genomes to library")
282
+
283
+ cmd = f"find {genomes_dir} -name '*.gz' | xargs -n 1 -P {threads} gunzip -k"
284
+ run_cmd(cmd)
285
+
286
+ cmd = f"find {genomes_dir} -name '*.gbff'"
287
+ files = run_cmd(cmd, return_output=True)
288
+ for file in files:
289
+ if os.path.exists(f"{file}.fna"):
290
+ continue
291
+ cmd = f"any2fasta -u {file} > {file}.fna"
292
+ run_cmd(cmd)
293
+
294
+ cmd = f"find {genomes_dir} -type f -name '*.fna'"
295
+ files = run_cmd(cmd, return_output=True)
296
+ logger.info(f"Found {len(files)} genomes to add to {db_name} library")
297
+ else:
298
+ organisms = DB_TYPE_CONFIG.get(db_type, [db_type])
299
+ files = []
300
+ for organism in organisms:
301
+ cmd = f"find {cache_dir}/refseq/{organism} -name '*.fna'"
302
+ org_files = run_cmd(cmd, return_output=True)
303
+ logger.info(f"Found {len(org_files)} genomes for {organism}")
304
+ files.extend(org_files)
305
+
306
+ return files
307
+
308
+
309
+ def save_md5_file(*args, **kwargs):
310
+ global md5_file
311
+ with open(md5_file, "w") as out_file:
312
+ for line in hashes:
313
+ out_file.write(line + "\n")
314
+ logger.info(f"Saved {len(hashes)} md5 hashes")
315
+
316
+
317
+ def add_to_library(
318
+ cache_dir, cwd, genomes_dir, db_type, db_name,
319
+ limit, batch_size, threads, use_k2
320
+ ):
321
+ os.chdir(cwd)
322
+ os.makedirs(cwd / db_name / "library", exist_ok=True)
323
+
324
+ files = get_files(genomes_dir, cache_dir, db_type, db_name, threads)
325
+ if limit:
326
+ logger.info(f"Limiting number of genomes to {limit}")
327
+ files = files[:limit]
328
+
329
+ step = batch_size
330
+ dynamic_step = len(files) // 10
331
+ step = min(step, dynamic_step)
332
+ if step == 0:
333
+ step = 1
334
+
335
+ logger.info(f"Using step size of {step}")
336
+
337
+ file_count = len(files)
338
+ start = datetime.datetime.now()
339
+
340
+ if use_k2:
341
+ for index, file in enumerate(files, start=1):
342
+ if index % step == 0:
343
+ duration = datetime.datetime.now() - start
344
+ average_speed = duration / step
345
+ eta = (file_count - index) * average_speed
346
+ logger.info(f"{datetime.datetime.now()}: Added {index} genomes in {duration}. ETA: {eta}")
347
+ start = datetime.datetime.now()
348
+
349
+ cmd = f"k2 add-to-library --db {db_name} --files {file}"
350
+ run_cmd(cmd, no_output=True)
351
+
352
+ logger.info(f"Added downloaded genomes to library")
353
+ end = datetime.datetime.now()
354
+ print(f"Time taken: {end - start}")
355
+ return
356
+
357
+ global hashes
358
+ global md5_file
359
+ md5_file = cwd / db_name / "library" / "added.md5"
360
+
361
+ if os.path.exists(md5_file):
362
+ with open(md5_file, "r") as in_file:
363
+ hashes = {line.strip() for line in in_file}
364
+
365
+ logger.info(f"Found {len(hashes)} md5 hashes in {md5_file}")
366
+
367
+ for index, file in enumerate(files, start=1):
368
+ if index % step == 0:
369
+ duration = datetime.datetime.now() - start
370
+ average_speed = duration / step
371
+ eta = (file_count - index) * average_speed
372
+ logger.info(f"{datetime.datetime.now()}: Added {index} genomes in {duration}. ETA: {eta}")
373
+ start = datetime.datetime.now()
374
+
375
+ if not os.path.exists(f"{file}.md5"):
376
+ md5sum = hash_file(file)
377
+ with open(f"{file}.md5", "w") as fh:
378
+ fh.write(md5sum)
379
+ else:
380
+ with open(f"{file}.md5", "r") as in_file:
381
+ md5sum = in_file.read()
382
+
383
+ if md5sum in hashes:
384
+ continue
385
+
386
+ cmd = f"kraken2-build --db {db_name} --add-to-library {file} --threads {threads}"
387
+ run_cmd(cmd, no_output=True)
388
+
389
+ with open(md5_file, "a") as out_file:
390
+ out_file.write(md5sum + "\n")
391
+
392
+ hashes.add(md5sum)
393
+
394
+ end = datetime.datetime.now()
395
+ print(f"Time taken: {end - start}")
396
+
397
+ logger.info(f"Added downloaded genomes to library")
398
+
399
+
400
+ @click.group(no_args_is_help=True, epilog=f"Config file: {get_config_path()}")
401
+ def cli():
402
+ pass
403
+
404
+
405
+ @cli.command(no_args_is_help=True, context_settings={"ignore_unknown_options": True})
406
+ @click.option('--tool', default='kraken2', type=click.Choice(list(REQUIRED_BINS)), help='Classifier to build the database for')
407
+ @click.option('--db-type', default=None, help='database type to build')
408
+ @click.option('--db-name', default=None, help='database name to build')
409
+ @click.option('--genomes-dir', default=None, help='Directory containing genomes')
410
+ @click.option('--cache-dir', default=create_cache_dir(), help='Cache directory')
411
+ @click.option('--threads', default=multiprocessing.cpu_count(), help='Number of threads to use', type=int)
412
+ @click.option('--load-factor', default=0.7, help='Proportion of the hash table to be populated. Used only for kraken2')
413
+ @click.option('--kmer-len', default=35, help='Kmer length in bp/aa. Used only in build task', type=int)
414
+ @click.option('--min-len', default=31, help='Minimizer/window length in bp/aa. Used only in build task', type=int)
415
+ @click.option('--level', default='leaves', type=click.Choice(['leaves', 'species', 'genus', 'assembly']), help='Taxonomic level to group sequences by. Used only for ganon2')
416
+ @click.option('--limit', default=None, help='Limit number of genomes to use', type=int)
417
+ @click.option('--batch-size', default=1000, help='Number of genomes to add to library at a time. Used only for kraken2', type=int)
418
+ @click.option('--force', is_flag=True, help='Force download and build')
419
+ @click.option('--rebuild', is_flag=True, help='Clean existing build files and re-build')
420
+ @click.option('--fast-build', is_flag=True, help='Non deterministic but faster build. Used only for kraken2')
421
+ @click.option('--use-k2', is_flag=True, help='Use k2 CLI instead of kraken2-build. Used only for kraken2')
422
+ @click.argument('ganon_args', nargs=-1, type=click.UNPROCESSED)
423
+ @click.pass_context
424
+ def build(
425
+ context,
426
+ tool: str, db_type: str, db_name, cache_dir, genomes_dir,
427
+ threads, load_factor, kmer_len: int, min_len, level: str, limit: int, batch_size: int,
428
+ force: bool, rebuild, fast_build: bool, use_k2: bool, ganon_args
429
+ ):
430
+ if tool in ('ganon', 'ganon2') and ganon_args:
431
+ logger.info(f"Passing through unrecognized options to native ganon build: {' '.join(ganon_args)}")
432
+ run_basic_checks(tool, use_k2)
433
+ cmd = "ganon build " + " ".join(shlex.quote(arg) for arg in ganon_args)
434
+ run_cmd(cmd)
435
+ return
436
+
437
+ logger.info(f"Building {tool} database of type {db_type}")
438
+ run_basic_checks(tool, use_k2)
439
+ cwd = Path(os.getcwd())
440
+
441
+ if cache_dir == '.':
442
+ cache_dir = cwd
443
+
444
+ if not db_name:
445
+ db_name = f"{tool}_{context.params['db_type']}"
446
+
447
+ if force:
448
+ if tool == 'kraken2':
449
+ run_cmd(f"rm -rf {db_name}")
450
+ run_cmd(f"mkdir -p {db_name}")
451
+ else:
452
+ run_cmd(f"rm -f {db_name}.*")
453
+
454
+ logger.info(f"Using cache directory {cache_dir}")
455
+
456
+ download_taxanomy(cache_dir)
457
+
458
+ if not genomes_dir:
459
+ download_genomes(cache_dir, cwd, db_type, db_name, threads, force)
460
+
461
+ if tool == 'kraken2':
462
+ add_to_library(
463
+ cache_dir, cwd, genomes_dir, db_type, db_name,
464
+ limit, batch_size, threads, use_k2
465
+ )
466
+ build_db(
467
+ cache_dir, cwd, db_type, db_name, threads, kmer_len, min_len,
468
+ fast_build, rebuild, load_factor, use_k2
469
+ )
470
+ elif tool == 'ganon2':
471
+ build_ganon2(
472
+ cache_dir, cwd, genomes_dir, db_type, db_name, threads,
473
+ kmer_len, min_len, level, rebuild
474
+ )
475
+
476
+
477
+ @cli.group(name='config', invoke_without_command=True)
478
+ @click.pass_context
479
+ def config(context):
480
+ if context.invoked_subcommand is not None:
481
+ return
482
+
483
+ parser = load_config()
484
+ items = parser.items(CONFIG_SECTION)
485
+ if not items:
486
+ logger.info(f"No config set. Config file: {get_config_path()}")
487
+ return
488
+
489
+ for key, value in items:
490
+ print(f"{key} = {value}")
491
+
492
+
493
+ @config.command(name='get')
494
+ @click.argument('key')
495
+ def config_get(key):
496
+ parser = load_config()
497
+ if not parser.has_option(CONFIG_SECTION, key):
498
+ logger.error(f"{key} not set")
499
+ sys.exit(1)
500
+ print(parser.get(CONFIG_SECTION, key))
501
+
502
+
503
+ @config.command(name='set')
504
+ @click.argument('key')
505
+ @click.argument('value')
506
+ def config_set(key, value):
507
+ parser = load_config()
508
+ parser.set(CONFIG_SECTION, key, value)
509
+ save_config(parser)
510
+ logger.info(f"Set {key} = {value}")
511
+
512
+
513
+ DOCTOR_BINS = ["ncbi-genome-download", "kraken2-build", "k2", "ganon", "any2fasta", "wget", "tar", "gunzip"]
514
+
515
+
516
+ @cli.command(name='doctor')
517
+ def doctor():
518
+ for binary in DOCTOR_BINS:
519
+ path = shutil.which(binary)
520
+ status = path if path else "MISSING"
521
+ print(f"{binary:<20} {status}")
522
+
523
+ print()
524
+ print(f"Cache dir: {create_cache_dir()}")
525
+ print(f"Config file: {get_config_path()}")
526
+
527
+
528
+ if __name__ == '__main__':
529
+ cli()