From 9121a4b080fec1e85832a1df752ab0aa0d02dd76 Mon Sep 17 00:00:00 2001 From: Ben Steadman Date: Tue, 26 Mar 2019 07:45:46 +0000 Subject: [PATCH] regionTrimmer.py Python3 refactor --- contrib/regionTrimmer.py | 217 +++++++++++++++++++++---------------- test/test_regionTrimmer.py | 58 ++++++++++ 2 files changed, 180 insertions(+), 95 deletions(-) create mode 100644 test/test_regionTrimmer.py diff --git a/contrib/regionTrimmer.py b/contrib/regionTrimmer.py index 7e26fce..98a7200 100644 --- a/contrib/regionTrimmer.py +++ b/contrib/regionTrimmer.py @@ -1,40 +1,51 @@ -#!/usr/bin/env python +#!/usr/bin/env python3 """Deletes outlying and unconnected regions""" +import argparse import logging -import os -import sys -import glob +from pathlib import Path import networkx logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) + def get_region_file_from_node(regionset_path, node): - return os.path.join(regionset_path, 'r.%d.%d.mca' % node) + return regionset_path / ('r.%d.%d.mca' % node) + def get_nodes(regionset_path): - return [tuple(map(int, r.split('.')[1:3])) \ - for r in glob.glob(os.path.join(regionset_path, 'r.*.*.mca'))] + return [ + tuple(int(x) for x in r.stem.split('.')[1:3]) + for r in regionset_path.glob('r.*.*.mca') + ] + def generate_edges(graph): offsets = (-1, 1) nodes = graph.nodes() for node in nodes: for offset in offsets: - graph.add_edges_from((node, offset_node) for offset_node in \ - [(node[0] + offset, node[1]), (node[0], node[1] + offset), \ - (node[0] + offset, node[1] + offset)] \ - if offset_node in nodes) + graph.add_edges_from( + (node, offset_node) + for offset_node in [ + (node[0] + offset, node[1]), + (node[0], node[1] + offset), + (node[0] + offset, node[1] + offset), + ] + if offset_node in nodes + ) return graph + def generate_subgraphs(nodes): graph = networkx.Graph() graph.add_nodes_from(nodes) generate_edges(graph) - return graph, networkx.connected_component_subgraphs(graph) + return graph, [graph.subgraph(c) for c in networkx.connected_components(graph)] + def get_graph_bounds(graph): nodes = graph.nodes() @@ -45,114 +56,130 @@ def get_graph_bounds(graph): min(n[1] for n in nodes), ) + def get_graph_center_by_bounds(bounds): dx = bounds[0] - bounds[1] dy = bounds[2] - bounds[3] - return (dx / 2 + bounds[1], dy / 2 + bounds[3]) + return (dx // 2 + bounds[1], dy // 2 + bounds[3]) -def main(*args, **options): - if len(args) < 1: - logger.error('Missing region directory argument') - return - for path in args: - logger.info('Using regionset path: %s', path) + +def trim_regions(graph, regions_path, dry_run=True, filter_func=lambda n: True): + regions = [ + (n, get_region_file_from_node(regions_path, n)) + for n in graph.nodes() + if filter_func(n) + ] + logger.info("Trimming regions: %s", ", ".join(x[1] for x in regions)) + for n, region_file in regions: + graph.remove_node(n) + if dry_run is False: + unlink_file(region_file) + + +def is_outside_main(center, main_section_bounds): + return center[0] <= main_section_bounds[0] and center[0] >= main_section_bounds[1] and \ + center[1] <= main_section_bounds[2] and center[1] >= main_section_bounds[3] + + +def is_outside_bounds(node, trim_center, trim_bounds): + return node[0] >= trim_center[0] + trim_bounds[0] or \ + node[0] <= trim_center[0] - trim_bounds[0] or \ + node[1] >= trim_center[1] + trim_bounds[1] or \ + node[1] <= trim_center[1] - trim_bounds[1] + +def unlink_file(path): + try: + path.unlink() + except OSError as err: + logger.warn("Unable to delete file: %s", path) + logger.warn("Error recieved was: %s", err) + + +def main(args): + for path in args.paths: + logger.info("Using regionset path: %s", path) nodes = get_nodes(path) if not len(nodes): - logger.error('Found no nodes, are you sure there are .mca files in %s ?', - path) + logger.error("Found no nodes, are you sure there are .mca files in %s ?", + path) return - logger.info('Found %d nodes', len(nodes)) - logger.info('Generating graphing nodes...') + logger.info("Found %d nodes", len(nodes)) + logger.info("Generating graphing nodes...") graph, subgraphs = generate_subgraphs(nodes) assert len(graph.nodes()) == sum(len(sg.nodes()) for sg in subgraphs) if len(subgraphs) == 1: - logger.warn('All regions are contiguous, the needful is done!') + logger.warn("All regions are contiguous, the needful is done!") return - logger.info('Found %d discrete region sections', len(subgraphs)) + logger.info("Found %d discrete region sections", len(subgraphs)) subgraphs = sorted(subgraphs, key=lambda sg: len(sg), reverse=True) for i, sg in enumerate(subgraphs): - logger.info('Region section #%02d: %04d nodes', i+1, len(sg.nodes())) + logger.info("Region section #%02d: %04d nodes", i + 1, len(sg.nodes())) bounds = get_graph_bounds(sg) - logger.info('Bounds: %d <-> %d x %d <-> %d', *get_graph_bounds(sg)) + logger.info("Bounds: %d <-> %d x %d <-> %d", *get_graph_bounds(sg)) center = get_graph_center_by_bounds(bounds) - logger.info('Center: %d x %d', *center) + logger.info("Center: %d x %d", *center) main_section = subgraphs[0] main_section_bounds = get_graph_bounds(main_section) main_section_center = get_graph_center_by_bounds(main_section_bounds) - logger.info('Using %d node graph as main section,', len(main_section.nodes())) + logger.info("Using %d node graph as main section,", len(main_section.nodes())) satellite_sections = subgraphs[1:] for ss in satellite_sections: bounds = get_graph_bounds(ss) center = get_graph_center_by_bounds(bounds) - logger.info('Checking satellite section with %d nodes, %d <-> %d x %d <-> %d bounds and %d x %d center', - len(ss.nodes()), *(bounds + center)) - if options['trim_disconnected']: - logger.info('Trimming regions: %s', ', '.join( - get_region_file_from_node(path, n) for n in ss.nodes())) - for n, region_file in ((n, get_region_file_from_node(path, n)) \ - for n in ss.nodes()): - ss.remove_node(n) - if not options['dry_run']: - unlink_file(region_file) - if options['trim_outside_main']: - if center[0] <= main_section_bounds[0] and center[0] >= main_section_bounds[1] and \ - center[1] <= main_section_bounds[2] and center[1] >= main_section_bounds[3]: - logger.info('Section falls inside main section bounds, ignoring') + logger.info(("Checking satellite section with %d nodes, " + "%d <-> %d x %d <-> %d bounds and %d x %d center"), + len(ss.nodes()), *(bounds + center)) + + if args.trim_disconnected: + trim_regions(ss, path, dry_run=args.dry_run) + + if args.trim_outside_main: + if is_outside_main(ss, center, main_section_bounds): + logger.info("Section is outside main section bounds") + trim_regions(ss, path, dry_run=args.dry_run) else: - logger.info('Section is outside main section bounds') - logger.info('Trimming regions: %s', ', '.join( - get_region_file_from_node(path, n) for n in ss.nodes())) - for n, region_file in ((n, get_region_file_from_node(path, n)) \ - for n in ss.nodes()): - ss.remove_node(n) - if not options['dry_run']: - unlink_file(region_file) - if options['trim_outside_bounds']: - x = map(int, options['trim_outside_bounds'].split(',')) - if len(x) == 4: - trim_center = x[:2] - trim_bounds = x[2:] - elif len(x) == 2: - trim_center = main_section_center - trim_bounds = x - else: - logger.error('Invalid center/bound value: %s', - options['trim_outside_bounds']) - continue - for node in ss.nodes(): - if node[0] >= trim_center[0] + trim_bounds[0] or \ - node[0] <= trim_center[0] - trim_bounds[0] or \ - node[1] >= trim_center[1] + trim_bounds[1] or \ - node[1] <= trim_center[1] - trim_bounds[1]: - region_file = get_region_file_from_node(path, node) - logger.info('Region falls outside specified bounds, trimming: %s', - region_file) - ss.remove_node(node) - if not options['dry_run']: - unlink_file(region_file) + logger.info("Section falls inside main section bounds, ignoring") -def unlink_file(path): - try: - os.unlink(path) - except OSError as err: - logger.warn('Unable to delete file: %s', path) - logger.warn('Error recieved was: %s', err) + if args.trim_outside_bounds: + logger.info("Checking regions outside specified bounds") + trim_center = args.trim_outside_bounds.get("center", main_section_center) + trim_bounds = args.trim_outside_bounds["bounds"] + trim_regions(ss, path, dry_run=args.dry_run, + filter_func=lambda n: is_outside_bounds(n, trim_center, trim_bounds)) -if __name__ == '__main__': - import optparse +def dir_path(path): + p = Path(path) + if not p.is_dir(): + raise argparse.ArgumentTypeError("Not a valid directory path") + return p + + +def center_bound(value): + x = [int(v) for v in value.split(",")] + if len(x) == 4: + return {"center": x[:2], "bounds": x[2:]} + elif len(x) == 2: + return {"bounds": x} + else: + raise argparse.ArgumentTypeError("Invalid center/bound value") + + +if __name__ == "__main__": logging.basicConfig() - parser = optparse.OptionParser( - usage='Usage: %prog [options] ') - parser.add_option('-D', '--trim-disconnected', action='store_true', default=False, - help='Trim all disconnected regions') - parser.add_option('-M', '--trim-outside-main', action='store_true', default=False, - help='Trim disconnected regions outside main section bounds') - parser.add_option('-B', '--trim-outside-bounds', default=False, - metavar='[center_X,center_Y,]bound_X,bound_Y', - help='Trim outside given bounds (given as [center_X,center_Y,]bound_X,bound_Y)') - parser.add_option('-n', '--dry-run', action='store_true', default=False, - help='Don\'t actually delete anything') - opts, args = parser.parse_args() - main(*args, **vars(opts)) + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("paths", metavar="", nargs="+", type=dir_path) + parser.add_argument("-D", "--trim-disconnected", action="store_true", + default=False, help="Trim all disconnected regions") + parser.add_argument("-M", "--trim-outside-main", action="store_true", + default=False, help="Trim disconnected regions outside main section bounds") + parser.add_argument("-B", "--trim-outside-bounds", + metavar="[center_X,center_Y,]bound_X,bound_Y", type=center_bound, + help=("Trim outside given bounds " + "(given as [center_X,center_Y,]bound_X,bound_Y)")) + parser.add_argument("-n", "--dry-run", action="store_true", default=False, + help="Don't actually delete anything") + args = parser.parse_args() + main(args) diff --git a/test/test_regionTrimmer.py b/test/test_regionTrimmer.py new file mode 100644 index 0000000..85800a4 --- /dev/null +++ b/test/test_regionTrimmer.py @@ -0,0 +1,58 @@ +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory + +import networkx + +import contrib.regionTrimmer as region_trimmer + + +class TestRegionTrimmer(unittest.TestCase): + def test_get_nodes(self): + coords = [(0, 0), (0, -1), (-1, 0), (-1, -1)] + with TemporaryDirectory() as tmpdirname: + region_file = Path(tmpdirname) + for x, z in coords: + region_fname = "r.{x}.{z}.mca".format(x=x, z=z) + (region_file / region_fname).touch() + + nodes = region_trimmer.get_nodes(region_file) + self.assertListEqual(sorted(nodes), sorted(coords)) + + def test_get_nodes_returns_empty_list_when_no_region_files(self): + with TemporaryDirectory() as tmpdirname: + region_file = Path(tmpdirname) + (region_file / "not_region_file.txt").touch() + nodes = region_trimmer.get_nodes(region_file) + self.assertListEqual(nodes, []) + + def test_get_region_file_from_node(self): + node = (0, 0) + regionset_path = Path('/path/to/regions') + + self.assertEqual(region_trimmer.get_region_file_from_node( + regionset_path, node), Path('/path/to/regions/r.0.0.mca')) + + def test_get_graph_bounds(self): + """ Should return (max_x, min_x, max_z, min_z) of all nodes + """ + graph = networkx.Graph() + graph.add_nodes_from([(0, 0), (0, -1), (-1, 0), (-1, -1)]) + + self.assertEqual(region_trimmer.get_graph_bounds(graph), (0, -1, 0, -1)) + + def test_get_graph_center_by_bounds(self): + self.assertEqual(region_trimmer.get_graph_center_by_bounds((0, -1, 0, -1)), (-1, -1)) + + def test_generate_edges(self): + graph = networkx.Graph() + graph.add_nodes_from([(0, 0), (0, -1), (-1, 0), (-1, -1)]) + + graph = region_trimmer.generate_edges(graph) + expected = [((-1, 0), (-1, -1)), + ((0, -1), (-1, -1)), + ((0, 0), (-1, -1)), + ((0, 0), (-1, 0)), + ((0, 0), (0, -1))] + + self.assertListEqual(sorted(list(graph.edges)), expected)