From 7b23f2796b6e99c25f3cca8f84990f439c4aeb38 Mon Sep 17 00:00:00 2001
From: Nikita Shulga <nikita.shulga@gmail.com>
Date: Sat, 30 Sep 2023 14:06:22 -0700
Subject: [PATCH] [BE] Add `s3_management` to the linted folders

---
 .lintrunner.toml              |  2 +-
 pyproject.toml                |  3 +++
 s3_management/backup_conda.py | 14 ++++++++------
 s3_management/manage.py       | 17 +++++++----------
 4 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/.lintrunner.toml b/.lintrunner.toml
index c551cb732..5c63a6c5d 100644
--- a/.lintrunner.toml
+++ b/.lintrunner.toml
@@ -2,7 +2,7 @@ merge_base_with = "origin/main"
 
 [[linter]]
 code = 'RUFF'
-include_patterns = ['test/smoke_test/*.py']
+include_patterns = ['test/smoke_test/*.py', 's3_management/*.py']
 command = [
     'python3',
     'tools/linter/adapters/ruff_linter.py',
diff --git a/pyproject.toml b/pyproject.toml
index efa884a07..b16714852 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,9 @@
 [tool.ruff]
 target-version = "py38"
 line-length = 120
+ignore = [
+    "PERF401",
+]
 select = [
     "B",
     "C4",
diff --git a/s3_management/backup_conda.py b/s3_management/backup_conda.py
index 06926589d..7dafa32b4 100644
--- a/s3_management/backup_conda.py
+++ b/s3_management/backup_conda.py
@@ -4,9 +4,9 @@
 # Do not use unless you know what you are doing
 # Usage:  python backup_conda.py --version 1.6.0
 
-import conda.api
 import boto3
 from typing import List, Optional
+import conda.api
 import urllib
 import os
 import hashlib
@@ -22,8 +22,11 @@ def compute_md5(path:str) -> str:
         return hashlib.md5(f.read()).hexdigest()
 
 
-def download_conda_package(package:str, version:Optional[str] = None, depends:Optional[str] = None, channel:Optional[str] = None) -> List[str]:
-    packages = conda.api.SubdirData.query_all(package, channels = [channel] if channel is not None else None, subdirs = _known_subdirs)
+def download_conda_package(package:str, version:Optional[str] = None,
+                           depends:Optional[str] = None, channel:Optional[str] = None) -> List[str]:
+    packages = conda.api.SubdirData.query_all(package,
+                                              channels = [channel] if channel is not None else None,
+                                              subdirs = _known_subdirs)
     rc = []
 
     for pkg in packages:
@@ -36,9 +39,8 @@ def download_conda_package(package:str, version:Optional[str] = None, depends:Op
         os.makedirs(pkg.subdir, exist_ok = True)
         fname = f"{pkg.subdir}/{pkg.fn}"
         if not os.path.exists(fname):
-            with open(fname, "wb") as f:
-                with urllib.request.urlopen(pkg.url) as url:
-                    f.write(url.read())
+            with open(fname, "wb") as f, urllib.request.urlopen(pkg.url) as url:
+                f.write(url.read())
         if compute_md5(fname) != pkg.md5:
             print(f"md5 of {fname} is {compute_md5(fname)} does not match {pkg.md5}")
             continue
diff --git a/s3_management/manage.py b/s3_management/manage.py
index 0df513d32..57af9259b 100644
--- a/s3_management/manage.py
+++ b/s3_management/manage.py
@@ -184,9 +184,7 @@ def nightly_packages_to_show(self: S3IndexType) -> Set[S3Object]:
             if package_name not in PACKAGE_ALLOW_LIST:
                 to_hide.add(obj)
                 continue
-            if packages[package_name] >= KEEP_THRESHOLD:
-                to_hide.add(obj)
-            elif between_bad_dates(package_build_time):
+            if packages[package_name] >= KEEP_THRESHOLD or between_bad_dates(package_build_time):
                 to_hide.add(obj)
             else:
                 packages[package_name] += 1
@@ -215,14 +213,13 @@ def gen_file_list(
         )
         subdir = self._resolve_subdir(subdir) + '/'
         for obj in objects:
-            if package_name is not None:
-                if self.obj_to_package_name(obj) != package_name:
-                    continue
+            if package_name is not None and self.obj_to_package_name(obj) != package_name:
+                continue
             if self.is_obj_at_root(obj) or obj.key.startswith(subdir):
                 yield obj
 
     def get_package_names(self, subdir: Optional[str] = None) -> List[str]:
-        return sorted(set(self.obj_to_package_name(obj) for obj in self.gen_file_list(subdir)))
+        return sorted({self.obj_to_package_name(obj) for obj in self.gen_file_list(subdir)})
 
     def normalize_package_version(self: S3IndexType, obj: S3Object) -> str:
         # removes the GPU specifier from the package name as well as
@@ -284,7 +281,7 @@ def to_simple_package_html(
         # Adding html footer
         out.append('  </body>')
         out.append('</html>')
-        out.append('<!--TIMESTAMP {}-->'.format(int(time.time())))
+        out.append(f'<!--TIMESTAMP {int(time.time())}-->')
         return '\n'.join(out)
 
     def to_simple_packages_html(
@@ -303,7 +300,7 @@ def to_simple_packages_html(
         # Adding html footer
         out.append('  </body>')
         out.append('</html>')
-        out.append('<!--TIMESTAMP {}-->'.format(int(time.time())))
+        out.append(f'<!--TIMESTAMP {int(time.time())}-->')
         return '\n'.join(out)
 
     def upload_legacy_html(self) -> None:
@@ -412,7 +409,7 @@ def main():
     args = parser.parse_args()
     action = "Saving" if args.do_not_upload else "Uploading"
     if args.prefix == 'all':
-        for prefix in PREFIXES_WITH_HTML.keys():
+        for prefix in PREFIXES_WITH_HTML:
             print(f"INFO: {action} indices for '{prefix}'")
             idx = S3Index.from_S3(prefix=prefix)
             if args.do_not_upload: