Add label-wrap feature (#16)

airflow_diagrams/cli.py

@@ -16,6 +16,7 @@
     load_mappings,
     render_jinja,
     to_var,
+    wrap_str,
 )
 
 app = Typer()
@@ -92,6 +93,12 @@ def generate(  # dead: disable
         "-v",
         help="Verbose output i.e. useful for debugging purposes.",
     ),
+    label_wrap: Optional[str] = Option(
+        None,
+        "--label-wrap",
+        "-lw",
+        help="Specify either a number for label width or a separator to indicate when to wrap a label.",
+    ),
 ) -> None:
     if verbose:
         echo("💬 Running with verbose output..")
@@ -133,7 +140,9 @@ def generate(  # dead: disable
                 diagram_nodes.append(
                     dict(
                         task_var=to_var(airflow_task.task_id),
-                        task_id=airflow_task.task_id,
+                        task_id=wrap_str(airflow_task.task_id, label_wrap)
+                        if label_wrap
+                        else airflow_task.task_id,
                         class_name=match_class_ref.class_name,
                     ),
                 )

airflow_diagrams/utils.py

@@ -2,6 +2,7 @@
 from hashlib import md5
 from os.path import dirname
 from pathlib import Path
+from textwrap import wrap
 
 import yaml
 from jinja2.environment import Environment
@@ -66,3 +67,16 @@ def to_var(string: str) -> str:
     :returns: a valid python variable name.
     """
     return f"_{md5(string.encode()).hexdigest()}"  # nosec
+
+
+def wrap_str(string: str, indicator: str) -> str:
+    """
+    Wrap a string with newline chars based on indicator.
+
+    :params string: The string to wrap.
+    :params indicator: Specify either a number for width or a separator to indicate when to wrap.
+    :returns: the wrapped string.
+    """
+    if indicator.isdigit():
+        return "\\n".join(wrap(string, int(indicator), break_on_hyphens=False))
+    return "\\n".join(string.split(indicator))

tests/test_utils.py

@@ -1,4 +1,4 @@
-from airflow_diagrams.utils import load_abbreviations, load_mappings, to_var
+from airflow_diagrams.utils import load_abbreviations, load_mappings, to_var, wrap_str
 
 
 def test_load_abbreviations(mocker):
@@ -23,3 +23,13 @@ def test_to_var():
     var = to_var("foo-bar")
     assert "-" not in var and "." not in var
     assert var.startswith("_")
+
+
+def test_wrap_str__indicator_number():
+    """Test wrapping a string based on number"""
+    assert wrap_str("foobarfoobar", "10") == "foobarfoob\\nar"
+
+
+def test_wrap_str__indicator_separator():
+    """Test wrapping a string based on separator"""
+    assert wrap_str("foobarfoobar", "foo") == "\\nbar\\nbar"