Ascii art first

metagenopolis · Mar 13, 2024 · 757ff99 · 757ff99
1 parent 1a3a77b
commit 757ff99
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 24 deletions.
diff --git a/meteor/meteor.py b/meteor/meteor.py
@@ -593,27 +593,27 @@ def get_arguments() -> Namespace:  # pragma: no cover
         type=str,
         help="Gap character (default -).",
     )
+    tree_parser.add_argument(
+        "-f",
+        dest="format",
+        default="txt",
+        choices=["png", "svg", "pdf", "txt"],
+        type=str,
+        help="Output image format (default txt).",
+    )
     tree_parser.add_argument(
         "-w",
         dest="width",
-        default=200,
+        default=500,
         type=int,
-        help="Output image width (default 200px).",
+        help="Output image width (default 500px).",
     )
     tree_parser.add_argument(
         "-H",
         dest="height",
-        default=200,
+        default=500,
         type=int,
-        help="Output image height (default 200px).",
-    )
-    tree_parser.add_argument(
-        "-f",
-        dest="format",
-        default="png",
-        choices=["png", "svg", "pdf"],
-        type=str,
-        help="Output image format (default png).",
+        help="Output image height (default 500px).",
     )
     tree_parser.add_argument(
         "-o",

diff --git a/meteor/treebuilder.py b/meteor/treebuilder.py
@@ -18,7 +18,7 @@
 from typing import Type, Dict, List
 from tempfile import mkdtemp
 import ete3  # type: ignore[import]
-from ete3 import Tree  # , TreeStyle
+from ete3 import Tree, TreeStyle
 from meteor.phylogeny import Phylogeny
 import logging
 import sys
@@ -92,19 +92,23 @@ def execute(self) -> None:
         # Analyze tree data
         for msp_file in msp_file_list:
             tree_file = self.meteor.tree_dir / f"{msp_file.stem}.tree"
-            # img_file = self.meteor.tree_dir / f"{msp_file.stem}.{self.format}"
+            img_file = self.meteor.tree_dir / f"{msp_file.stem}.{self.format}"
             msp_tree = Tree(str(tree_file.resolve()))
             # Generate a distance msp by msp
             matrix = self.get_msp_distance(msp_tree)
             matrix.to_csv(self.meteor.tree_dir / f"{msp_file.stem}.tsv", sep="\t")
             # Draw trees
-            # ts = TreeStyle()
-            # ts.show_leaf_name = True
-            # ts.show_branch_length = True
-            # msp_tree.render(
-            #     str(img_file.resolve()),
-            #     w=self.width,
-            #     h=self.height,
-            #     units="px",
-            #     dpi=300,
-            # )
+            if self.format == "txt":
+                with img_file.open("wt", encoding="UTF-8") as outfile:
+                    outfile.write(msp_tree.get_ascii(show_internal=True))
+            else:
+                ts = TreeStyle()
+                ts.show_leaf_name = True
+                ts.show_branch_length = True
+                msp_tree.render(
+                    str(img_file.resolve()),
+                    w=self.width,
+                    h=self.height,
+                    units="px",
+                    dpi=300,
+                )