Merge pull request #11 from interline-io/extract-osmconvert

Planet extractor: add osmctools, script output

planetutils/osm_planet_extract.py

@@ -13,12 +13,19 @@ def main():
     parser.add_argument('--name', help='Name to give to extract file.')
     parser.add_argument('--bbox', help='Bounding box for extract file. Format for coordinates: left,bottom,right,top')
     parser.add_argument('--verbose', help="Verbose output", action='store_true')
+    parser.add_argument('--toolchain', help='OSM toolchain', default='osmosis')
+    parser.add_argument('--commands', help='Output a command list instead of performing action, e.g. for parallel usage', action='store_true')
     args = parser.parse_args()
 
     if args.verbose:
         log.set_verbose()
 
-    p = Planet(args.osmpath)
+    if args.toolchain == 'osmosis':
+        p = PlanetExtractorOsmosis(args.osmpath)
+    elif args.toolchain == 'osmctools':
+        p = PlanetExtractorOsmconvert(args.osmpath)
+    else:
+        parser.error('unknown toolchain: %s'%args.toolchain)
 
     bboxes = {}
     if args.csv:
@@ -30,7 +37,12 @@ def main():
     else:
         parser.error('must specify --csv, --geojson, or --bbox and --name')
 
-    p.extract_bboxes(bboxes, outpath=args.outpath)
+    if args.commands:
+        commands = p.extract_commands(bboxes, outpath=args.outpath)
+        for i in commands:
+            print " ".join(i)
+    else:
+        p.extract_bboxes(bboxes, outpath=args.outpath)
 
 if __name__ == '__main__':
     main()

planetutils/planet.py

@@ -15,21 +15,18 @@ def __init__(self, osmpath=None, grain='hour', changeset_url=None, osmosis_workd
         d, p = os.path.split(osmpath)
         self.osmosis_workdir = osmosis_workdir or os.path.join(d, '%s.workdir'%p)
 
-    def osmosis(self, *args):
-        cmd = ['osmosis'] + list(args)
-        log.debug(' '.join(cmd))
+    def command(self, args):
+        log.debug(args)
         return subprocess.check_output(
-            cmd,
+            args,
             shell=False
         )
 
+    def osmosis(self, *args):
+        return self.command(['osmosis'] + list(args))
+
     def osmconvert(self, *args):
-        cmd = ['osmconvert'] + list(args)
-        log.debug(' '.join(cmd))
-        return subprocess.check_output(
-            cmd,
-            shell=False
-        )
+        return self.command(['osmconvert'] + list(args))
 
     def get_timestamp(self):
         timestamp = self.osmconvert(
@@ -48,12 +45,23 @@ def get_timestamp(self):
             ][0]
         return timestamp.strip()
 
-    def download_planet(self):
-        raise NotImplementedError
+class Planet(PlanetBase):
+    pass
 
-    def update_planet(self, outpath, grain='hour', changeset_url=None):
+class PlanetExtractor(PlanetBase):
+    def extract_bboxes(self, bboxes, workers=1, outpath='.'):
         raise NotImplementedError
 
+    def extract_bbox(self, name, bbox, workers=1, outpath='.'):
+        return self.extract_bboxes({name: bbox}, outpath=outpath, workers=workers)
+
+    def extract_commands(self, bboxes, outpath='.'):
+        args = []
+        self.command = lambda x:args.append(x)
+        self.extract_bboxes(bboxes, outpath=outpath)
+        return args
+
+class PlanetExtractorOsmosis(PlanetExtractor):
     def extract_bboxes(self, bboxes, workers=1, outpath='.'):
         args = []
         args += ['--read-pbf-fast', self.osmpath, 'workers=%s'%int(workers)]
@@ -73,11 +81,24 @@ def extract_bboxes(self, bboxes, workers=1, outpath='.'):
             args += arg
         self.osmosis(*args)
 
+class PlanetExtractorOsmconvert(PlanetExtractor):
+    def extract_bboxes(self, bboxes, workers=1, outpath='.'):
+        for name, bbox in bboxes.items():
+            self.extract_bbox(name, bbox, outpath=outpath)
+
     def extract_bbox(self, name, bbox, workers=1, outpath='.'):
-        return self.extract_bboxes({name: bbox}, outpath=outpath, workers=workers)
+        validate_bbox(bbox)
+        left, bottom, right, top = bbox
+        args = [
+            self.osmpath,
+            '-b=%s,%s,%s,%s'%(left, bottom, right, top),
+            '-o=%s'%os.path.join(outpath, '%s.osm.pbf'%name)
+        ]
+        self.osmconvert(*args)
 
-class Planet(PlanetBase):
-    pass
+class PlanetDownloader(PlanetBase):
+    def download_planet(self):
+        raise NotImplementedError
 
 class PlanetDownloaderHttp(PlanetBase):
     def _download(self, url, outpath):
@@ -124,6 +145,11 @@ def _get_planets(self, bucket, prefix, match):
                 objs.append(obj)
         return objs
 
+
+class PlanetUpdater(PlanetBase):
+    def update_planet(self, outpath, grain='hour', changeset_url=None):
+        raise NotImplementedError
+
 class PlanetUpdaterOsmupdate(PlanetBase):
     pass
 

tests/test_planet.py

@@ -8,6 +8,9 @@
 TESTFILE_TIMESTAMP = '2018-02-02T22:34:43Z'
 TEST_BBOX = [-122.430439,37.766508,-122.379670,37.800052]
 
+# import planetutils.log as log
+# log.set_verbose()
+
 class TestPlanetBase(unittest.TestCase):
     def test_osmosis(self):
         p = planet.PlanetBase(TESTFILE)
@@ -27,11 +30,13 @@ def test_get_timestamp(self):
         p = planet.PlanetBase(TESTFILE)
         self.assertEquals(p.get_timestamp(), TESTFILE_TIMESTAMP)
 
-    def test_extract_bbox(self, name=None, bbox=None):
-        name = name or 'test'
-        bbox = bbox or TEST_BBOX
+class TestPlanetExtractor(unittest.TestCase):
+    kls = None
+    def extract_bbox(self):
+        name = 'test'
+        bbox = TEST_BBOX
         d = tempfile.mkdtemp()
-        p = planet.PlanetBase(TESTFILE)
+        p = self.kls(TESTFILE)
         outfile = os.path.join(d, '%s.osm.pbf'%name)
         p.extract_bbox(name, bbox, outpath=d)
         self.assertTrue(os.path.exists(outfile))
@@ -40,6 +45,16 @@ def test_extract_bbox(self, name=None, bbox=None):
         os.unlink(outfile)
         os.rmdir(d)
 
+class TestPlanetExtractorOsmconvert(TestPlanetExtractor):
+    kls = planet.PlanetExtractorOsmconvert
+    def test_extract_bbox(self):
+        self.extract_bbox()
+
+class TestPlanetExtractorOsmosis(TestPlanetExtractor):
+    kls = planet.PlanetExtractorOsmosis
+    def test_extract_bbox(self):
+        self.extract_bbox()
+
 class TestPlanetDownloaderHttp(unittest.TestCase):
     def test_download_planet(self):
         p = planet.PlanetDownloaderHttp('test.osm.pbf')