Merge pull request #185 from KernelTuner/output_file_writer

Moved .json extension checking to function and added module description to avoid code smells

Author: benvanwerkhoven

kernel_tuner/file_utils.py

@@ -1,3 +1,5 @@
+""" This module contains utility functions for operations on files, mostly JSON cache files """
+
 import os
 import json
 import subprocess
@@ -6,8 +8,6 @@
 from importlib.metadata import requires, version, PackageNotFoundError
 from packaging.requirements import Requirement
 
-from jsonschema import validate
-
 from kernel_tuner import util
 
 schema_dir = os.path.dirname(os.path.realpath(__file__)) + "/schema"
@@ -45,6 +45,13 @@ def get_configuration_validity(objective) -> str:
     return errorstring
 
 
+def filename_ensure_json_extension(filename: str) -> str:
+    """ Check if the filename has a .json extension, if not, add it """
+    if filename[-5:] != ".json":
+        filename += ".json"
+    return filename
+
+
 def store_output_file(output_filename, results, tune_params, objective="time"):
     """ Store the obtained auto-tuning results in a JSON output file
 
@@ -63,8 +70,7 @@ def store_output_file(output_filename, results, tune_params, objective="time"):
     :type objective: string
 
     """
-    if output_filename[-5:] != ".json":
-        output_filename += ".json"
+    output_filename = filename_ensure_json_extension(output_filename)
 
     timing_keys = [
         "compile_time", "benchmark_time", "framework_time", "strategy_time",
@@ -171,8 +177,7 @@ def store_metadata_file(metadata_filename, target="nvidia"):
     :type target: string
 
     """
-    if metadata_filename[-5:] != ".json":
-        metadata_filename += ".json"
+    metadata_filename = filename_ensure_json_extension(metadata_filename)
     metadata = {}
 
     # lshw only works on Linux, this intentionally raises a FileNotFoundError when ran on systems that do not have it

test/test_file_utils.py

@@ -1,9 +1,11 @@
-from kernel_tuner.file_utils import store_output_file, store_metadata_file, output_file_schema, validate
+from kernel_tuner.file_utils import store_output_file, store_metadata_file, output_file_schema
 from kernel_tuner.util import delete_temp_file
 from .test_integration import fake_results
 from .test_runners import env, cache_filename, tune_kernel
+
 import pytest
 import json
+from jsonschema import validate
 
 
 def test_store_output_file(env):