Merge pull request #54 from jurraca/use-pathlib

context: use pathlib instead of os.path

kartograf/context.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-import os
+from pathlib import Path
 import sys
 import time
 
@@ -20,12 +20,12 @@ def __init__(self, args):
 
         self.reproduce = self.args.reproduce is not None
         if self.reproduce:
-            items = os.listdir(self.args.reproduce)
+            items = Path.iterdir(Path(self.args.reproduce))
             source_folders = []
             for folder in items:
-                full_path = os.path.join(self.args.reproduce, folder)
-                if os.path.isdir(full_path):
-                    source_folders.append(folder)
+                full_path = self.args.reproduce / folder
+                if Path.is_dir(full_path):
+                    source_folders.append(folder.name)
             # We override the args because we are reproducing and only where we
             # have data we try to use is, the actual args passed don't matter.
             self.args.irr = 'irr' in source_folders
@@ -44,47 +44,52 @@ def __init__(self, args):
             self.epoch_dir = str(int(utc_time_now))
             self.epoch_datetime = datetime.utcfromtimestamp(int(utc_time_now))
 
-        cwd = os.getcwd()
+        cwd = Path.cwd()
         # Data dir
         if self.reproduce:
             if not self.args.reproduce.endswith('/'):
                 self.args.reproduce += '/'
             self.data_dir = self.args.reproduce
         else:
-            self.data_dir = f"{cwd}/data/{self.epoch_dir}/"
-        self.data_dir_irr = f"{self.data_dir}irr/"
-        self.data_dir_rpki_cache = f"{self.data_dir}rpki/cache/"
-        self.data_dir_rpki_tals = f"{self.data_dir}rpki/tals/"
-        self.data_dir_collectors = f"{self.data_dir}collectors/"
+            self.data_dir = str(cwd / "data" / self.epoch_dir)
+
+        if Path(self.data_dir).exists() and not self.reproduce:
+            print("Not so fast, a folder with that epoch already exists.")
+            sys.exit()
+
+        self.data_dir_irr = str(Path(self.data_dir) / "irr")
+        self.data_dir_rpki_cache = str(Path(self.data_dir) / "rpki" / "cache")
+        self.data_dir_rpki_tals = str(Path(self.data_dir) / "rpki" / "tals")
+        self.data_dir_collectors = str(Path(self.data_dir) / "collectors")
         # Out dir
-        self.out_dir = f"{cwd}/out/{self.epoch_dir}/"
-        self.out_dir_irr = f"{self.out_dir}irr/"
-        self.out_dir_rpki = f"{self.out_dir}rpki/"
-        self.out_dir_collectors = f"{self.out_dir}collectors/"
+        self.out_dir = str(cwd / "out" / self.epoch_dir)
+        self.out_dir_irr = str(Path(self.out_dir) / "irr")
+        self.out_dir_rpki = str(Path(self.out_dir) / "rpki")
+        self.out_dir_collectors = str(Path(self.out_dir) / "collectors")
 
-        if os.path.exists(self.data_dir) and not self.reproduce:
+        if Path(self.data_dir).exists() and not self.reproduce:
             print("Not so fast, a folder with that epoch already exists.")
             sys.exit()
 
         # We skip creating the folders if we are reproducing a run.
         if not self.reproduce:
-            os.makedirs(self.data_dir_rpki_cache)
-            os.makedirs(self.data_dir_rpki_tals)
+            Path(self.data_dir_rpki_cache).mkdir(parents=True)
+            Path(self.data_dir_rpki_tals).mkdir(parents=True)
             if self.args.irr:
-                os.makedirs(self.data_dir_irr)
+                Path(self.data_dir_irr).mkdir(parents=True)
             if self.args.routeviews:
-                os.makedirs(self.data_dir_collectors)
-        os.makedirs(self.out_dir_rpki)
+                Path(self.data_dir_collectors).mkdir(parents=True)
+        Path(self.out_dir_rpki).mkdir(parents=True)
         if self.args.irr:
-            os.makedirs(self.out_dir_irr)
+            Path(self.out_dir_irr).mkdir(parents=True)
         if self.args.routeviews:
-            os.makedirs(self.out_dir_collectors)
+            Path(self.out_dir_collectors).mkdir(parents=True)
 
-        self.final_result_file = f"{self.out_dir}final_result.txt"
+        self.final_result_file = str(Path(self.out_dir) / "final_result.txt")
 
         self.max_encode = self.args.max_encode
 
         if self.args.debug:
-            self.debug_log = f"{self.out_dir}debug.log"
+            self.debug_log = str(Path(self.out_dir) / "debug.log")
         else:
             self.debug_log = ""

kartograf/rpki/parse.py

@@ -1,4 +1,5 @@
 import json
+from pathlib import Path
 from typing import Dict
 
 from kartograf.bogon import (
@@ -12,8 +13,8 @@
 
 @timed
 def parse_rpki(context):
-    raw_input = f"{context.out_dir_rpki}rpki_raw.json"
-    rpki_res = f"{context.out_dir_rpki}rpki_final.txt"
+    raw_input = Path(context.out_dir_rpki) / "rpki_raw.json"
+    rpki_res = Path(context.out_dir_rpki) / "rpki_final.txt"
 
     output_cache: Dict[str, [str, str]] = {}
 

tests/test_context.py

@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 import pytest
 from kartograf.context import Context
 from kartograf.cli import create_parser
@@ -21,7 +22,7 @@ def test_basic_map_context(parser, tmp_path):
     assert isinstance(context.epoch, str)
     assert isinstance(int(context.epoch), int)
     assert context.max_encode == 33521664
-    assert context.debug_log.endswith('debug.log')
+    assert Path(context.debug_log).name == 'debug.log'
 
 def test_map_context_with_reproduce(parser, tmp_path):
     # Setup a mock reproduction directory
@@ -53,10 +54,17 @@ def test_directory_creation(parser, tmp_path):
     os.chdir(tmp_path)
     context = Context(args)
 
-    assert os.path.exists(context.data_dir_rpki_cache)
-    assert os.path.exists(context.data_dir_rpki_tals)
-    assert os.path.exists(context.data_dir_irr)
-    assert os.path.exists(context.data_dir_collectors)
-    assert os.path.exists(context.out_dir_rpki)
-    assert os.path.exists(context.out_dir_irr)
-    assert os.path.exists(context.out_dir_collectors)
+    assert isinstance(context.data_dir_rpki_cache, str)
+    assert Path(context.data_dir_rpki_cache).exists()
+    assert isinstance(context.data_dir_rpki_tals, str)
+    assert Path(context.data_dir_rpki_tals).exists()
+    assert isinstance(context.data_dir_irr, str)
+    assert Path(context.data_dir_irr).exists()
+    assert isinstance(context.data_dir_collectors, str)
+    assert Path(context.data_dir_collectors).exists()
+    assert isinstance(context.out_dir_rpki, str)
+    assert Path(context.out_dir_rpki).exists()
+    assert isinstance(context.out_dir_irr, str)
+    assert Path(context.out_dir_irr).exists()
+    assert isinstance(context.out_dir_collectors, str)
+    assert Path(context.out_dir_collectors).exists()

tests/test_merge.py

@@ -40,10 +40,10 @@ def test_merge_from_fixtures(tmp_path):
     and validates against expected networks,
     i.e. subnets are merged into the root network appropriately.
     '''
-    testdir = Path(__file__).parent
-    base_nets, base_results = __read_test_vectors(testdir / "data/base_file.csv")
+    data_dir = Path(__file__).parent / "data"
+    base_nets, base_results = __read_test_vectors(data_dir / "base_file.csv")
     base_path = tmp_path / "base.txt"
-    extra_nets, extra_results = __read_test_vectors(testdir / "data/extra_file.csv")
+    extra_nets, extra_results = __read_test_vectors(data_dir / "extra_file.csv")
     extra_path = tmp_path / "extra.txt"
     # write the networks to disk, generating ASNs for each network
     generate_ip_file(base_path, build_file_lines(base_nets, generate_asns(len(base_nets))))