Skip to content

Commit 052789b

Browse files
eaton-labeaton-lab
committed
cli updates
1 parent 8cdeb3f commit 052789b

File tree

6 files changed

+593
-130
lines changed

6 files changed

+593
-130
lines changed

‎toytree/cli/cli_draw.py‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
#!/usr/bin/env python
2+
3+
from pathlib import Path
4+
import sys
5+
import textwrap
6+
import tempfile
7+
import os
8+
import platform
9+
import subprocess
10+
import shutil
11+
import webbrowser
12+
from argparse import ArgumentParser, RawDescriptionHelpFormatter
13+
from .make_wide import make_wide
14+
from loguru import logger
15+
16+
17+
KWARGS = dict(
18+
prog="draw",
19+
usage="draw [options]",
20+
help="return ...",
21+
formatter_class=make_wide(RawDescriptionHelpFormatter, 120, 140),
22+
description=textwrap.dedent("""
23+
-------------------------------------------------------------------
24+
| draw: generate tree drawing as ascii or ...
25+
-------------------------------------------------------------------
26+
| The draw method generates a tree drawing in a variety of formats
27+
| and with many styling options.
28+
-------------------------------------------------------------------
29+
"""),
30+
epilog=textwrap.dedent("""
31+
Examples
32+
--------
33+
$ draw -i TRE.nwk -a
34+
$ draw -i TRE.nwk -f png -v
35+
$ draw -i TRE.nwk -f html -v -k ...
36+
$ draw -i TRE.nwk -f pdf -o /tmp/DRAWING
37+
$ draw -i TRE.nwk -v -N fill=red -E stroke=pink -T font-size=10px
38+
$ root -i TRE.nwk -n R | draw -i - -v
39+
""")
40+
)
41+
42+
43+
def get_parser_draw(parser: ArgumentParser | None = None) -> ArgumentParser:
44+
"""Return a parser tool for this method.
45+
"""
46+
# create parser or connect as subparser to cli parser
47+
if parser:
48+
KWARGS['name'] = KWARGS.pop("prog")
49+
parser = parser.add_parser(**KWARGS)
50+
else:
51+
KWARGS.pop("help")
52+
parser = ArgumentParser(**KWARGS)
53+
54+
# path args
55+
parser.add_argument("-i", "--input", type=Path, metavar="path", required=True, help="input tree file (nwk, nex or nhx)")
56+
parser.add_argument("-o", "--output", type=Path, metavar="path", help="optional basename of outfile path. If None prints to STDOUT")
57+
# option
58+
parser.add_argument("-a", "--ascii", action="store_true", help="print ascii tree (overrides other draw args)")
59+
parser.add_argument("-f", "--format", choices=["html", "svg", "pdf", "png"], default="png", help="file format of drawing [png]")
60+
parser.add_argument("-v", "--view", type=str, metavar="app", const="auto", nargs="?", help="open drawing in default viewer, or provide an app name")
61+
parser.add_argument("-e", "--ladderize", action="store_true", help="ladderize the tree")
62+
parser.add_argument("-k", "--kwargs", type=str, metavar="str", nargs="*", help="any supported toytree.draw kwargs as 'key=val'")
63+
parser.add_argument("-I", "--internal-labels", type=str, metavar="str", help="parse internal node feature (e.g., support) [auto]")
64+
parser.add_argument("-l", "--log-level", type=str, metavar="level", default="INFO", help="stderr logging level (DEBUG, [INFO], WARNING, ERROR)")
65+
66+
parser.add_argument("-ts", "--tree-style", type=str, metavar="str", help="base tree style [[None], 'r', 'c', 's', 'o', 'b']")
67+
parser.add_argument("-N", "--node-style", type=str, metavar="str", nargs="+", help="node style args")
68+
parser.add_argument("-E", "--edge-style", type=str, metavar="str", nargs="+", help="edge style args")
69+
parser.add_argument("-T", "--tip-labels-style", type=str, metavar="str", nargs="+", help="tip labels style args")
70+
parser.add_argument("-ns", "--node-sizes", type=int, metavar="int", nargs="+", default=[6], help="node sizes")
71+
parser.add_argument("-nc", "--node-colors", type=str, metavar="str", nargs="+", default=["#262626"], help="node colors")
72+
parser.add_argument("-tl", "--tip-labels-align", type=bool, metavar="bool", nargs="+", help="align tip labels")
73+
# parser.add_argument("-L", "--log-file", type=Path, metavar="path", help="append stderr log to a file")
74+
return parser
75+
76+
77+
78+
def open_with_default_viewer(path: str) -> bool:
79+
"""Try to open a file with the system's default app.
80+
Returns True on (likely) success, False if we had no good option.
81+
"""
82+
path = os.path.abspath(path)
83+
system = platform.system()
84+
85+
try:
86+
if system == "Windows":
87+
# Uses the file association in the registry
88+
os.startfile(path) # type: ignore[attr-defined]
89+
return True
90+
91+
elif system == "Darwin": # macOS
92+
subprocess.Popen(["open", path])
93+
return True
94+
95+
else:
96+
# Most Linux/Unix desktops support xdg-open
97+
opener = (
98+
shutil.which("xdg-open")
99+
or shutil.which("gio")
100+
or shutil.which("gnome-open")
101+
or shutil.which("kde-open")
102+
)
103+
if opener:
104+
subprocess.Popen([opener, path])
105+
return True
106+
107+
# Fallback: try the webbrowser module
108+
file_url = f"file://{path}"
109+
if webbrowser.open(file_url):
110+
return True
111+
112+
except Exception:
113+
# You might want to log this if you have logging set up
114+
logger.bind(name="toytree").error("could not find a default viewer to open drawing file")
115+
return False
116+
117+
118+
119+
def run_draw(args):
120+
from toytree import save
121+
from toytree.io.src.treeio import tree
122+
# from toytree.utils.src.logger_setup import set_log_level
123+
# set_log_level(args.log_level)
124+
125+
# parse the tree
126+
if args.input == Path("-"):
127+
data = sys.stdin.read()
128+
tre = tree(data, internal_labels=args.internal_labels)
129+
else:
130+
data = args.input.expanduser().absolute()
131+
tre = tree(data, internal_labels=args.internal_labels)
132+
133+
# ascii tree drawing
134+
if args.ascii:
135+
print(tre.treenode.draw_ascii(), sys.stdout)
136+
return 0
137+
138+
# create drawing
139+
# if args.kwargs:
140+
# print(args.kwargs)
141+
# kwargs = dict(tuple(i.split("=")) for i in args.kwargs)
142+
# else:
143+
# kwargs = {}
144+
canvas, axes, mark = tre.draw(
145+
tree_style=args.tree_style,
146+
node_style=dict(tuple(i.split("=") for i in args.node_style)) if args.node_style else {},
147+
edge_style=dict(tuple(i.split("=") for i in args.edge_style)) if args.edge_style else {},
148+
tip_labels_style=dict(tuple(i.split("=") for i in args.tip_labels_style)) if args.tip_labels_style else {},
149+
node_sizes=args.node_sizes if len(args.node_sizes) > 1 else args.node_sizes[0],
150+
# node_colors=args.node_colors if len(args.node_colors) > 1 else args.node_colors[0],
151+
# tip_labels_align=args.tip_labels_align,
152+
# **kwargs
153+
)
154+
canvas.style["background-color"] = "white"
155+
156+
# write file to tmp or named file
157+
suffix = "." + args.format.lower()
158+
if args.output:
159+
prefix = Path(args.output)
160+
if not prefix.name:
161+
prefix = prefix / "toytree"
162+
out = Path(f"{prefix}").with_suffix(suffix)
163+
else:
164+
out = tempfile.NamedTemporaryFile(prefix="toytree", suffix=suffix, delete=False)
165+
out = out.name
166+
save(canvas, out)
167+
168+
# optionally view the file
169+
if args.view:
170+
open_with_default_viewer(out)
171+
172+
173+
def main():
174+
parser = get_parser_draw()
175+
args = parser.parse_args()
176+
run_draw(args)
177+
178+
179+
if __name__ == "__main__":
180+
try:
181+
main()
182+
# except ToytreeError as exc:
183+
# logger.bind(name="toytree").error(exc)
184+
except Exception as exc:
185+
logger.bind(name="toytree").exception(exc)

‎toytree/cli/cli_get_node_data.py‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
#!/usr/bin/env python
2+
3+
"""
4+
"""
5+
6+
# from typing import List
7+
import sys
8+
import textwrap
9+
from pathlib import Path
10+
from argparse import ArgumentParser, RawDescriptionHelpFormatter
11+
from loguru import logger
12+
from .make_wide import make_wide
13+
# from ortholab.utils.parallel import run_pipeline # , run_with_pool
14+
# from ortholab.utils.logger_setup import set_log_level
15+
16+
17+
KWARGS = dict(
18+
prog="node-data",
19+
usage="node-data [options]",
20+
help="return table of feature data stored in a tree file",
21+
formatter_class=make_wide(RawDescriptionHelpFormatter, 120, 140),
22+
description=textwrap.dedent("""
23+
-------------------------------------------------------------------
24+
| node-data: return...
25+
-------------------------------------------------------------------
26+
| The node-data method ...
27+
-------------------------------------------------------------------
28+
"""),
29+
epilog=textwrap.dedent("""
30+
Examples
31+
--------
32+
$ node-data -i TRE.nwk -H
33+
$ node-data -i TRE.nwk -n A B C
34+
$ node-data -i TRE.nwk -n ~prefix -f name dist
35+
$ node-data -i TRE.nwk -s ',' > DATA.csv
36+
$ cat TRE.nwk | get-node-data -i -
37+
""")
38+
)
39+
40+
41+
def string_or_stdin_parse(intree: str) -> str:
42+
"""If TREE is stdin then return the string from stdin."""
43+
if intree == "-":
44+
return sys.stdin.read().strip()
45+
return intree
46+
47+
48+
def get_parser_get_node_data(parser: ArgumentParser | None = None) -> ArgumentParser:
49+
"""Return a parser tool for this method.
50+
"""
51+
# create parser or connect as subparser to cli parser
52+
if parser:
53+
KWARGS['name'] = KWARGS.pop("prog")
54+
parser = parser.add_parser(**KWARGS)
55+
else:
56+
KWARGS.pop("help")
57+
parser = ArgumentParser(**KWARGS)
58+
59+
# path args
60+
parser.add_argument("-i", "--input", type=string_or_stdin_parse, metavar="path", required=True, help="input CDS sequence (aligned or unaligned)")
61+
parser.add_argument("-o", "--output", type=Path, metavar="path", help="optional outfile path name. If None prints to STDOUT")
62+
parser.add_argument("-n", "--nodes", type=str, metavar="str", nargs="*", help="one or more names or regular expressions to select nodes")
63+
parser.add_argument("-f", "--features", type=str, metavar="str", nargs="*", help="subselect one or more features to return")
64+
parser.add_argument("-t", "--tips-only", action="store_true", help="only return tip (leaf) node data")
65+
# options
66+
parser.add_argument("-s", "--separator", type=str, metavar="str", default="\t", help=r"separator character [\t]")
67+
parser.add_argument("-m", "--missing", type=str, metavar="str", help=r"a value to set for missing data to replace NaN")
68+
parser.add_argument("-H", "--human-readable", action="store_true", help="use easy to read variable column spacing. Overrides -s")
69+
parser.add_argument("-I", "--internal-labels", type=str, metavar="str", help="internal node features (e.g., support) [auto]")
70+
parser.add_argument("-l", "--log-level", type=str, metavar="level", default="INFO", help="stderr logging level (DEBUG, [INFO], WARNING, ERROR)")
71+
# parser.add_argument("-f", "--force", action="store_true", help="overwrite existing result files in outdir")
72+
# parser.add_argument("-L", "--log-file", type=Path, metavar="path", help="append stderr log to a file")
73+
return parser
74+
75+
76+
def run_get_node_data(args):
77+
from toytree.io.src.treeio import tree
78+
79+
# parse the tree
80+
if args.input == Path("-"):
81+
data = sys.stdin.read()
82+
tre = tree(data, internal_labels=args.internal_labels)
83+
else:
84+
data = args.input.expanduser().absolute()
85+
tre = tree(data, internal_labels=args.internal_labels)
86+
87+
# get data
88+
if args.tips_only:
89+
data = tre.get_tip_data(args.features)
90+
else:
91+
data = tre.get_node_data(args.features)
92+
93+
# subset nodes
94+
# TODO: support casting str back to int for node selections?
95+
if args.nodes:
96+
nodes = tre.get_nodes(*args.nodes)
97+
data = data.loc[[i.idx for i in nodes], :]
98+
99+
# print
100+
out = sys.stdout if args.output is None else args.output
101+
if args.human_readable:
102+
datastr = data.to_string()
103+
sys.stdout.write(f"{datastr}\n")
104+
else:
105+
data.to_csv(out, sep=args.separator)
106+
107+
108+
def main():
109+
parser = get_parser_get_node_data()
110+
args = parser.parse_args()
111+
run_get_node_data(args)
112+
113+
114+
if __name__ == "__main__":
115+
116+
try:
117+
main()
118+
except Exception as exc:
119+
logger.bind(name="toytree").error(exc)
120+

0 commit comments

Comments
 (0)