Skip to content

New --color-scheme argument #82

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 27 additions & 20 deletions code2flow/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand All @@ -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 = <
<table cellspacing="0" cellpadding="0" border="1"><tr><td>Code2flow Legend</td></tr><tr><td>
<table cellspacing="0">
<tr><td>Regular function</td><td width="50px" bgcolor='%s'></td></tr>
<tr><td>Trunk function (nothing calls this)</td><td bgcolor='%s'></td></tr>
<tr><td>Leaf function (this calls nothing else)</td><td bgcolor='%s'></td></tr>
<tr><td>Function call</td><td><font color='black'>&#8594;</font></td></tr>
</table></td></tr></table>
>];
}""" % (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 = <
<table cellspacing="0" cellpadding="0" border="1"><tr><td>Code2flow Legend</td></tr><tr><td>
<table cellspacing="0">
<tr><td>Regular function</td><td width="50px" bgcolor='%s'></td></tr>
<tr><td>Trunk function (nothing calls this)</td><td bgcolor='%s'></td></tr>
<tr><td>Leaf function (this calls nothing else)</td><td bgcolor='%s'></td></tr>
<tr><td>Function call</td><td><font color='black'>&#8594;</font></td></tr>
</table></td></tr></table>
>];
}""" % colors
return legend


LANGUAGES = {
Expand Down Expand Up @@ -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

Expand All @@ -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:
Expand Down Expand Up @@ -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.
Expand All @@ -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()
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -872,4 +878,5 @@ def main(sys_argv=None):
lang_params=lang_params,
subset_params=subset_params,
level=level,
color_scheme=args.color_scheme,
)
21 changes: 14 additions & 7 deletions code2flow/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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():
Expand Down
27 changes: 26 additions & 1 deletion tests/test_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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