Fix lint (#232)

Author: pankajroark

truss/tests/test_truss_handle.py

@@ -427,21 +427,29 @@ def test_update_requirements(custom_model_truss_dir_with_pre_and_post):
     sc_requirements = th.spec.requirements
     assert sc_requirements == requirements
 
+
 def test_add_python_requirements_new(custom_model_truss_dir_with_pre_and_post):
     th = TrussHandle(custom_model_truss_dir_with_pre_and_post)
     python_package_to_add = "tensorflow==2.3.1"
     th.add_python_requirement(python_package_to_add)
     sc_requirements = th.spec.requirements
     assert python_package_to_add in sc_requirements
 
-def test_add_python_requirements_version_change(custom_model_truss_dir_with_pre_and_post):
+
+def test_add_python_requirements_version_change(
+    custom_model_truss_dir_with_pre_and_post,
+):
     th = TrussHandle(custom_model_truss_dir_with_pre_and_post)
     th.add_python_requirement("tensorflow==2.3.1")
-    
+
     python_package_to_change_version = "tensorflow==2.3.5"
     th.add_python_requirement(python_package_to_change_version)
     sc_requirements = th.spec.requirements
-    assert python_package_to_change_version.split('==')[1]==sc_requirements[0].split('==')[1]
+    assert (
+        python_package_to_change_version.split("==")[1]
+        == sc_requirements[0].split("==")[1]
+    )
+
 
 def test_update_requirements_from_file(
     custom_model_truss_dir_with_pre_and_post, tmp_path

truss/truss_handle.py

@@ -80,10 +80,6 @@
     logger.addHandler(logging.StreamHandler(sys.stdout))
 
 
-def _python_req_name(python_requirement: str) -> str:
-    return pkg_resources.Requirement.parse(python_requirement).name
-
-
 class TrussHandle:
     def __init__(self, truss_dir: Path, validate: bool = True) -> None:
         self._truss_dir = truss_dir
@@ -405,8 +401,9 @@ def add_python_requirement(self, python_requirement: str):
         # Parse the added python requirements
         input_python_req_name = _python_req_name(python_requirement)
         new_reqs = [
-            req for req in self._spec.config.requirements
-                if _python_req_name(req) != input_python_req_name
+            req
+            for req in self._spec.config.requirements
+            if _python_req_name(req) != input_python_req_name
         ]
         new_reqs.append(python_requirement)
         self._update_config(lambda conf: replace(conf, requirements=new_reqs))
@@ -1114,3 +1111,8 @@ def _docker_image_from_labels(labels: Dict):
     images = get_images(labels)
     if images and isinstance(images, list):
         return images[0]
+
+
+def _python_req_name(python_requirement: str) -> str:
+    req = pkg_resources.Requirement.parse(python_requirement)
+    return req.name  # type: ignore