diff --git a/code2flow/engine.py b/code2flow/engine.py index ad13a43..0d43b06 100644 --- a/code2flow/engine.py +++ b/code2flow/engine.py @@ -11,7 +11,7 @@ from .javascript import Javascript from .ruby import Ruby from .php import PHP -from .model import (TRUNK_COLOR, LEAF_COLOR, NODE_COLOR, GROUP_TYPE, OWNER_CONST, +from .model import (COLOR_SCHEMES, GROUP_TYPE, OWNER_CONST, Edge, Group, Node, Variable, is_installed, flatten) VERSION = '2.5.1' @@ -24,19 +24,22 @@ "See the README at https://github.com/scottrogowski/code2flow." -LEGEND = """subgraph legend{ - rank = min; - label = "legend"; - Legend [shape=none, margin=0, label = < -
Code2flow Legend
- - - - - -
Regular function
Trunk function (nothing calls this)
Leaf function (this calls nothing else)
Function call→
- >]; -}""" % (NODE_COLOR, TRUNK_COLOR, LEAF_COLOR) +def _generate_legend(color_scheme='default'): + colors = COLOR_SCHEMES[color_scheme] + legend = """subgraph legend{ + rank = min; + label = "legend"; + Legend [shape=none, margin=0, label = < +
Code2flow Legend
+ + + + + +
Regular function
Trunk function (nothing calls this)
Leaf function (this calls nothing else)
Function call→
+ >]; + }""" % colors + return legend LANGUAGES = { @@ -229,7 +232,7 @@ def generate_json(nodes, edges): def write_file(outfile, nodes, edges, groups, hide_legend=False, - no_grouping=False, as_json=False): + no_grouping=False, as_json=False, color_scheme='default'): ''' Write a dot file that can be read by graphviz @@ -253,9 +256,9 @@ def write_file(outfile, nodes, edges, groups, hide_legend=False, content += f'splines="{splines}";\n' content += 'rankdir="LR";\n' if not hide_legend: - content += LEGEND + content += _generate_legend(color_scheme=color_scheme) for node in nodes: - content += node.to_dot() + ';\n' + content += node.to_dot(color_scheme=color_scheme) + ';\n' for edge in edges: content += edge.to_dot() + ';\n' if not no_grouping: @@ -672,7 +675,7 @@ def code2flow(raw_source_paths, output_file, language=None, hide_legend=True, exclude_namespaces=None, exclude_functions=None, include_only_namespaces=None, include_only_functions=None, no_grouping=False, no_trimming=False, skip_parse_errors=False, - lang_params=None, subset_params=None, level=logging.INFO): + lang_params=None, subset_params=None, color_scheme='default', level=logging.INFO): """ Top-level function. Generate a diagram based on source code. Can generate either a dotfile or an image. @@ -691,6 +694,7 @@ def code2flow(raw_source_paths, output_file, language=None, hide_legend=True, :param lang_params LanguageParams: Object to store lang-specific params :param subset_params SubsetParams: Object to store subset-specific params :param int level: logging level + :param str color_scheme: legend color scheme :rtype: None """ start_time = time.time() @@ -751,11 +755,11 @@ def code2flow(raw_source_paths, output_file, language=None, hide_legend=True, as_json = output_ext == 'json' write_file(fh, nodes=all_nodes, edges=edges, groups=file_groups, hide_legend=hide_legend, - no_grouping=no_grouping, as_json=as_json) + no_grouping=no_grouping, as_json=as_json, color_scheme=color_scheme) else: write_file(output_file, nodes=all_nodes, edges=edges, groups=file_groups, hide_legend=hide_legend, - no_grouping=no_grouping) + no_grouping=no_grouping, color_scheme=color_scheme) logging.info("Wrote output file %r with %d nodes and %d edges.", output_file, len(all_nodes), len(edges)) @@ -837,6 +841,8 @@ def main(sys_argv=None): help='add more logging') parser.add_argument( '--version', action='version', version='%(prog)s ' + VERSION) + parser.add_argument('--color-scheme', action='store', choices=list(COLOR_SCHEMES.keys()), default='default', + help="select a color for the graph legend") sys_argv = sys_argv or sys.argv[1:] args = parser.parse_args(sys_argv) @@ -872,4 +878,5 @@ def main(sys_argv=None): lang_params=lang_params, subset_params=subset_params, level=level, + color_scheme=args.color_scheme, ) diff --git a/code2flow/model.py b/code2flow/model.py index adb4a36..7623c00 100644 --- a/code2flow/model.py +++ b/code2flow/model.py @@ -2,11 +2,18 @@ import os -TRUNK_COLOR = '#966F33' -LEAF_COLOR = '#6db33f' EDGE_COLORS = ["#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7"] -NODE_COLOR = "#cccccc" + +# (NODE_COLOR, TRUNK_COLOR, LEAF_COLOR) +COLOR_SCHEMES = { + 'default': ('#cccccc', '#966F33', '#6db33f'), + 'blue': ('#138bfa', '#989a97', '#044a9a'), + 'desert': ('#fda653', '#ffecd5', '#d57720'), + 'pink': ('#b99cd6', '#ddd5d7', '#cb3870'), + 'light_blue': ('#ace2ff', '#feffff', '#5bcce0'), + 'aqua': ('#1a9e85', '#304960', '#1fdec3'), +} class Namespace(dict): @@ -401,7 +408,7 @@ def resolve_variables(self, file_groups): else: assert isinstance(variable.points_to, (Node, Group)) - def to_dot(self): + def to_dot(self, color_scheme='default'): """ Output for graphviz (.dot) files :rtype: str @@ -411,12 +418,12 @@ def to_dot(self): 'name': self.name(), 'shape': "rect", 'style': 'rounded,filled', - 'fillcolor': NODE_COLOR, + 'fillcolor': COLOR_SCHEMES[color_scheme][0], } if self.is_trunk: - attributes['fillcolor'] = TRUNK_COLOR + attributes['fillcolor'] = COLOR_SCHEMES[color_scheme][1] elif self.is_leaf: - attributes['fillcolor'] = LEAF_COLOR + attributes['fillcolor'] = COLOR_SCHEMES[color_scheme][2] ret = self.uid + ' [' for k, v in attributes.items(): diff --git a/tests/test_interface.py b/tests/test_interface.py index ad9541f..e13c9e3 100644 --- a/tests/test_interface.py +++ b/tests/test_interface.py @@ -13,6 +13,8 @@ from code2flow import model IMG_PATH = '/tmp/code2flow/output.png' +GRAPHVIZ_PATH = '/tmp/code2flow/out.gv' + if os.path.exists("/tmp/code2flow"): try: shutil.rmtree('/tmp/code2flow') @@ -95,7 +97,8 @@ def test_json(): assert jobj['graph']['directed'] is True assert isinstance(jobj['graph']['nodes'], dict) assert len(jobj['graph']['nodes']) == 4 - assert set(n['name'] for n in jobj['graph']['nodes'].values()) == {'simple_b::a', 'simple_b::(global)', 'simple_b::c.d', 'simple_b::b'} + assert set(n['name'] for n in jobj['graph']['nodes'].values()) == {'simple_b::a', 'simple_b::(global)', + 'simple_b::c.d', 'simple_b::b'} assert isinstance(jobj['graph']['edges'], list) assert len(jobj['graph']['edges']) == 4 @@ -202,6 +205,7 @@ def test_cli_log_quiet(mocker): logging.basicConfig.assert_called_once_with(format="Code2Flow: %(message)s", level=logging.WARNING) + def test_subset_cli(mocker): with pytest.raises(AssertionError): SubsetParams.generate(target_function='', upstream_depth=1, downstream_depth=0) @@ -221,3 +225,24 @@ def test_subset_cli(mocker): main(['test_code/py/subset_find_exception/two.py', '--target-function', 'func', '--upstream-depth', '1']) +def test_color_scheme_invalid_input(capsys): + with pytest.raises(SystemExit): + main(['test_code/py/simple_a', '--color-scheme', 'invalid']) + assert 'error: argument --color-scheme: invalid choice:' in capsys.readouterr().err + + +def test_color_scheme_none_given(capsys): + code2flow(['test_code/py/simple_a'], output_file=IMG_PATH) + assert os.path.exists(IMG_PATH) + + +def test_color_scheme_test_not_default(capsys): + code2flow(['test_code/py/two_file_simple', '--color-scheme', 'blue'], output_file=GRAPHVIZ_PATH, color_scheme='blue') + assert os.path.exists(GRAPHVIZ_PATH) + contents = '' + with open(GRAPHVIZ_PATH, 'r') as file: + contents = "".join(file.readlines()) + blue = ('#138bfa', '#989a97', '#044a9a') + assert blue[0] in contents + assert blue[1] in contents + assert blue[2] in contents