Skip to content

Commit d5fc947

Browse files
authored
Fix lint (#232)
1 parent 1a55494 commit d5fc947

File tree

2 files changed

+19
-9
lines changed

2 files changed

+19
-9
lines changed

truss/tests/test_truss_handle.py

+11-3
Original file line numberDiff line numberDiff line change
@@ -427,21 +427,29 @@ def test_update_requirements(custom_model_truss_dir_with_pre_and_post):
427427
sc_requirements = th.spec.requirements
428428
assert sc_requirements == requirements
429429

430+
430431
def test_add_python_requirements_new(custom_model_truss_dir_with_pre_and_post):
431432
th = TrussHandle(custom_model_truss_dir_with_pre_and_post)
432433
python_package_to_add = "tensorflow==2.3.1"
433434
th.add_python_requirement(python_package_to_add)
434435
sc_requirements = th.spec.requirements
435436
assert python_package_to_add in sc_requirements
436437

437-
def test_add_python_requirements_version_change(custom_model_truss_dir_with_pre_and_post):
438+
439+
def test_add_python_requirements_version_change(
440+
custom_model_truss_dir_with_pre_and_post,
441+
):
438442
th = TrussHandle(custom_model_truss_dir_with_pre_and_post)
439443
th.add_python_requirement("tensorflow==2.3.1")
440-
444+
441445
python_package_to_change_version = "tensorflow==2.3.5"
442446
th.add_python_requirement(python_package_to_change_version)
443447
sc_requirements = th.spec.requirements
444-
assert python_package_to_change_version.split('==')[1]==sc_requirements[0].split('==')[1]
448+
assert (
449+
python_package_to_change_version.split("==")[1]
450+
== sc_requirements[0].split("==")[1]
451+
)
452+
445453

446454
def test_update_requirements_from_file(
447455
custom_model_truss_dir_with_pre_and_post, tmp_path

truss/truss_handle.py

+8-6
Original file line numberDiff line numberDiff line change
@@ -80,10 +80,6 @@
8080
logger.addHandler(logging.StreamHandler(sys.stdout))
8181

8282

83-
def _python_req_name(python_requirement: str) -> str:
84-
return pkg_resources.Requirement.parse(python_requirement).name
85-
86-
8783
class TrussHandle:
8884
def __init__(self, truss_dir: Path, validate: bool = True) -> None:
8985
self._truss_dir = truss_dir
@@ -405,8 +401,9 @@ def add_python_requirement(self, python_requirement: str):
405401
# Parse the added python requirements
406402
input_python_req_name = _python_req_name(python_requirement)
407403
new_reqs = [
408-
req for req in self._spec.config.requirements
409-
if _python_req_name(req) != input_python_req_name
404+
req
405+
for req in self._spec.config.requirements
406+
if _python_req_name(req) != input_python_req_name
410407
]
411408
new_reqs.append(python_requirement)
412409
self._update_config(lambda conf: replace(conf, requirements=new_reqs))
@@ -1114,3 +1111,8 @@ def _docker_image_from_labels(labels: Dict):
11141111
images = get_images(labels)
11151112
if images and isinstance(images, list):
11161113
return images[0]
1114+
1115+
1116+
def _python_req_name(python_requirement: str) -> str:
1117+
req = pkg_resources.Requirement.parse(python_requirement)
1118+
return req.name # type: ignore

0 commit comments

Comments
 (0)