Add ruff rules for pyupgrade (#546)

Author: cbornet

libs/knowledge-graph/ragstack_knowledge_graph/knowledge_schema.py

@@ -96,12 +96,10 @@ def validate_graph_document(self, document: GraphDocument):
                 e.add_note(f"No edge type '{r.edge_type}")
             else:
                 relationship = next(
-                    (
-                        candidate
-                        for candidate in relationships
-                        if r.source_type in candidate.source_types
-                        if r.target_type in candidate.target_types
-                    )
+                    candidate
+                    for candidate in relationships
+                    if r.source_type in candidate.source_types
+                    if r.target_type in candidate.target_types
                 )
                 if relationship is None:
                     e.add_note(

libs/knowledge-graph/ragstack_knowledge_graph/traverse.py

@@ -190,7 +190,7 @@ def fetch_relationships(distance: int, source: Node) -> None:
             return results
 
 
-class AsyncPagedQuery(object):
+class AsyncPagedQuery:
     def __init__(self, depth: int, response_future: ResponseFuture):
         self.loop = asyncio.get_running_loop()
         self.depth = depth

libs/knowledge-graph/tests/conftest.py

@@ -35,10 +35,8 @@ def cassandra_port(db_keyspace: str) -> Iterator[int]:
     cassandra.start()
     wait_for_logs(cassandra, "Startup complete")
     cassandra.get_wrapped_container().exec_run(
-        (
-            f"""cqlsh -e "CREATE KEYSPACE {db_keyspace} WITH replication = """
-            '''{'class': 'SimpleStrategy', 'replication_factor': '1'};"'''
-        )
+        f"""cqlsh -e "CREATE KEYSPACE {db_keyspace} WITH replication = """
+        '''{'class': 'SimpleStrategy', 'replication_factor': '1'};"'''
     )
     port = cassandra.get_exposed_port(9042)
     print(f"Cassandra started. Port is {port}")

libs/knowledge-store/ragstack_knowledge_store/_mmr_helper.py

@@ -32,7 +32,7 @@ def update_redundancy(self, new_weighted_redundancy: float):
             self.score = self.weighted_similarity - self.weighted_redundancy
 
 
-class MmrHelper(object):
+class MmrHelper:
     dimensions: int
     """Dimensions of the embedding."""
 

libs/langchain/ragstack_langchain/graph_store/base.py

@@ -422,7 +422,7 @@ async def asearch(
                 "'mmr' or 'traversal'."
             )
 
-    def as_retriever(self, **kwargs: Any) -> "GraphStoreRetriever":
+    def as_retriever(self, **kwargs: Any) -> GraphStoreRetriever:
         """Return GraphStoreRetriever initialized from this GraphStore.
 
         Args:

libs/langchain/ragstack_langchain/graph_store/extractors/link_extractor.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Generic, Iterable, Set, TypeVar
+from typing import Generic, Iterable, TypeVar
 
 from ragstack_langchain.graph_store.links import Link
 
@@ -14,7 +14,7 @@ class LinkExtractor(ABC, Generic[InputT]):
     """Interface for extracting links (incoming, outgoing, bidirectional)."""
 
     @abstractmethod
-    def extract_one(self, input: InputT) -> Set[Link]:
+    def extract_one(self, input: InputT) -> set[Link]:
         """Add edges from each `input` to the corresponding documents.
 
         Args:

libs/ragulate/ragstack_ragulate/config/config_parser.py

@@ -28,7 +28,7 @@ def get_config(self) -> Config:
 
     @classmethod
     def from_file(cls, file_path: str) -> "ConfigParser":
-        with open(file_path, "r") as file:
+        with open(file_path) as file:
             config = yaml.safe_load(file)
 
             version = config.get("version", 0.1)

libs/tests-utils/ragstack_tests_utils/cassandra_container.py

@@ -9,7 +9,7 @@ def __init__(
         port: int = 9042,
         **kwargs,
     ) -> None:
-        super(CassandraContainer, self).__init__(image=image, **kwargs)
+        super().__init__(image=image, **kwargs)
         self.port = port
 
         self.with_exposed_ports(self.port)

libs/tests-utils/ragstack_tests_utils/test_data.py

@@ -13,12 +13,12 @@ def _get_test_data_path(file_name: str) -> str:
 
     @staticmethod
     def _get_text_file(file_name: str) -> str:
-        with open(TestData._get_test_data_path(file_name), "r") as f:
+        with open(TestData._get_test_data_path(file_name)) as f:
             return f.read()
 
     @staticmethod
     def _get_csv_embedding(csv_file_name: str) -> Embedding:
-        with open(TestData._get_test_data_path(csv_file_name), "r") as f:
+        with open(TestData._get_test_data_path(csv_file_name)) as f:
             reader = csv.reader(f)
             return [[float(value) for value in row] for row in reader]
 

pyproject.toml

@@ -62,9 +62,16 @@ select = [
     "E",
     "F",
     "I",
+    "UP",
     "W",
 ]
 
+[tool.ruff.lint.per-file-ignores]
+"**/libs/langchain/*" = [
+    "UP006", # Incompatible with Pydantic v1
+    "UP007", # Incompatible with Pydantic v1
+]
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

scripts/format-example-notebooks.py

@@ -16,7 +16,7 @@ def main():
             file = os.path.join(root, file)
             if ".ipynb_checkpoints" not in file and file.endswith(".ipynb"):
                 print("Formatting file: ", file)
-                with open(file, "r") as f:
+                with open(file) as f:
                     contents = f.read()
                     as_json = json.loads(contents)
                     cells = as_json["cells"]

scripts/generate-testspace-report.py

@@ -166,7 +166,7 @@ def parse_snyk_report(input_file: str):
         snykfile = os.path.join(input_file, snykfile)
         print("Reading file: " + snykfile)
 
-        with open(snykfile, "r") as file:
+        with open(snykfile) as file:
             data = json.load(file)
             for vulnerability in data.get("vulnerabilities", []):
                 title = vulnerability.get("title", "?")

scripts/parse-snyk-report.py

@@ -8,7 +8,7 @@
 input_file = sys.argv[1]
 output_file = sys.argv[2]
 
-with open(input_file, "r") as file:
+with open(input_file) as file:
     data = json.load(file)
     vulnerabilities = data.get("vulnerabilities", [])