Added ty LSP support for this repository, and auth_encrypt fix for lack of associated data

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arg-toolkit"
-version = "2.3.0"
+version = "2.3.1"
 description = "A collection of ComfyUI nodes for cryptography, steganography, encryption, decryption, and tools needed to both make and solve ARGs (Alternative Reality Games) and secret messages."
 readme = "README.md"
 requires-python = ">=3.12"
@@ -34,6 +34,8 @@ dev = [
   "pre-commit",  # runs linting on commit
   "pytest",  # testing
   "ruff",  # lintingclass
+  "ty", # lsp
+  "uv",
 ]
 
 [build-system]
@@ -65,7 +67,7 @@ testpaths = [
 
 [tool.ruff]
 # extend-exclude = ["static", "ci/templates"]
-line-length = 140
+line-length = 320
 src = ["src", "tests"]
 target-version = "py312"
 
@@ -74,9 +76,13 @@ target-version = "py312"
 select = [
   "S102",  # exec-builtin
   "S307",  # eval-used
-  "W293",
+  "W293",  # Blank line with whitespace
   "F",  # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names.
   # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f
+  "NPY",  # NumPy-specific rules
+  "SIM",  # flake8-simplify rules, to simplify the code more
+  "B",  # flake8-bugbear rules, to catch patterns with potential bugs
+  "PLR",  # Pylint refactor rules
 ]
 
 [tool.ruff.lint.flake8-quotes]
@@ -84,3 +90,10 @@ inline-quotes = "double"
 
 [tool.setuptools.package-data]
 "*" = ["*.*"]
+
+[tool.ty.analysis]
+# Suppress errors for all `test` modules except `test.foo`
+allowed-unresolved-imports = ["test.**", "!test.foo", "folder_paths"]
+
+[tool.ty.terminal]
+output-format = "concise"

src/cryptography_primitives/auth_encrypt.py

@@ -111,13 +111,6 @@ class AESAuth:  # Since all AES-based authenticated encryption techniques are th
                         "tooltip": "A random value to use. Should be 12 bytes in size.",
                     },
                 ),
-                "associated_data": (
-                    "BYTESLIKE",
-                    {
-                        "forceInput": True,
-                        "tooltip": "Additional data that should be authenticated with the key, but does not need to be encrypted. Can be None",
-                    },
-                ),
                 "mode": (
                     "BOOLEAN",
                     {
@@ -145,6 +138,13 @@ class AESAuth:  # Since all AES-based authenticated encryption techniques are th
                         "tooltip": "For AES-CCM specifically, it allows a tag length to be specified. Normally, this defaults to 16, but can be lowered to 4. Unless you know what you're doing, DO NOT CHANGE FROM THE DEFAULTS.",
                     },
                 ),
+                "associated_data": (
+                    "BYTESLIKE",
+                    {
+                        "forceInput": True,
+                        "tooltip": "Additional data that should be authenticated with the key, but does not need to be encrypted. Can be None",
+                    },
+                ),
             },
         }
 
@@ -152,7 +152,7 @@ class AESAuth:  # Since all AES-based authenticated encryption techniques are th
     RETURN_NAMES = ("encrypted_txt",)
     FUNCTION = "aesauth"
 
-    def aesauth(self, text, key: bytes, nonce, associated_data, mode, aes_type, ccm_tag_length):
+    def aesauth(self, text: bytes, key: bytes, nonce: bytes, mode: bool, aes_type: str, ccm_tag_length: int, associated_data: bytes | None):
         cipher_engines = {
             "AES-GCM": aead.AESGCM,
             "AES-GCM-SIV": aead.AESGCMSIV,
@@ -160,7 +160,7 @@ class AESAuth:  # Since all AES-based authenticated encryption techniques are th
         }
         if aes_type == "AES-SIV":
             cipher = aead.AESSIV(key)
-            associated_data = [associated_data, nonce]
+            associated_data = [associated_data, nonce]  # ty: ignore[invalid-assignment]
         elif aes_type == "AES-CCM":
             cipher = aead.AESCCM(key, ccm_tag_length)
         elif aes_type in cipher_engines:
@@ -169,14 +169,14 @@ class AESAuth:  # Since all AES-based authenticated encryption techniques are th
             raise ValueError("Invalid AES type chosen. Perhaps the node is broken? Try making a new one, as this is normally impossible.")
         if mode:
             if aes_type == "AES-SIV":
-                message = cipher.encrypt(text, associated_data)
+                message = cipher.encrypt(text, associated_data)  # ty: ignore[missing-argument, invalid-argument-type]
             else:
-                message = cipher.encrypt(nonce, text, associated_data)
+                message = cipher.encrypt(nonce, text, associated_data)  # ty: ignore[invalid-argument-type, too-many-positional-arguments]
         else:
             if aes_type == "AES-SIV":
-                message = cipher.decrypt(text, associated_data)
+                message = cipher.decrypt(text, associated_data)  # ty: ignore[missing-argument, invalid-argument-type]
             else:
-                message = cipher.decrypt(nonce, text, associated_data)
+                message = cipher.decrypt(nonce, text, associated_data)  # ty: ignore[invalid-argument-type, too-many-positional-arguments]
         return (message,)
 
 

src/steganography.py

@@ -9,20 +9,7 @@ import torch
 from PIL import Image
 import numpy as np
 
-# folder_path.py solver, just so the complaints will stop.
-current = os.path.dirname(os.path.realpath(__file__))
-while True:
-    candidate = os.path.join(current, "folder_paths.py")
-    if os.path.exists(candidate):
-        target_dir = current
-        break
-    parent = os.path.dirname(current)
-    if parent == current:
-        raise FileNotFoundError("folder_paths.py not found in any parent directories")
-    current = parent
-sys.path.insert(0, target_dir)
-
-import folder_paths  # This import is handled by the solver above looking for it inside the parent directories. This will most likely hit the one from base ComfyUI, assuming no other nodes or .py file carry this namespace.
+import folder_paths  # This import is handled by ComfyUI. This is a very jank solution, but it works, so don't ask.
 
 # Every node is inspired heavily by their reference implementations in their GitHub repository, with changes made to best use PyTorch as possible as it's the main way ComfyUI stores data
 # For more information, check https://docs.comfy.org/custom-nodes/backend/images_and_masks#images
@@ -80,7 +67,7 @@ class Stegano_LSB_Encode:
                     ["None"]
                     + [
                         name
-                        for name, _ in inspect.getmembers(stegano.lsb.generators, inspect.isfunction)
+                        for name, _ in inspect.getmembers(stegano.lsb.generators, inspect.isfunction)  # ty: ignore[possibly-missing-submodule]
                         if name
                         not in [
                             "carmichael",
@@ -117,7 +104,7 @@ class Stegano_LSB_Encode:
             if generator_type == "None":
                 generator_type = None
             else:
-                generator_func = getattr(stegano.lsb.generators, generator_type)
+                generator_func = getattr(stegano.lsb.generators, generator_type)  # ty: ignore[invalid-argument-type] Handled by the None checker above it  # ty: ignore[possibly-missing-submodule]
             if generator_type in ["LFSR", "ackermann", "ackermann_naive"]:
                 generator_type = generator_func(m=m)
             elif generator_type in ["ackermann_fast", "ackermann_slow"]:
@@ -127,7 +114,7 @@ class Stegano_LSB_Encode:
                 temp_path = os.path.join(temp_dir, f"image_{i}.png")
                 img_pil.save(temp_path)
                 generator_type = generator_func(temp_path)
-            img_pil = stegano.lsb.hide(img_pil, message, generator_type, encoding=encoding)
+            img_pil = stegano.lsb.hide(img_pil, message, generator_type, encoding=encoding)  # ty: ignore[invalid-argument-type] If it works, it works. At the backend, Pillow's Image format is also bytes
             img_np_out = np.array(img_pil)
             img_tensor_out = torch.from_numpy(img_np_out).float() / 255.0
             output_images.append(img_tensor_out)
@@ -176,7 +163,7 @@ class Stegano_LSB_Decode:
                     ["None"]
                     + [
                         name
-                        for name, _ in inspect.getmembers(stegano.lsb.generators, inspect.isfunction)
+                        for name, _ in inspect.getmembers(stegano.lsb.generators, inspect.isfunction)  # ty: ignore[possibly-missing-submodule]
                         if name
                         not in [
                             "carmichael",
@@ -213,7 +200,7 @@ class Stegano_LSB_Decode:
             if generator_type == "None":
                 generator_type = None
             else:
-                generator_func = getattr(stegano.lsb.generators, generator_type)
+                generator_func = getattr(stegano.lsb.generators, generator_type)  # ty: ignore[invalid-argument-type, possibly-missing-submodule]
             if generator_type in ["LFSR", "ackermann", "ackermann_naive"]:
                 generator_type = generator_func(m=m)
             elif generator_type in ["ackermann_fast", "ackermann_slow"]:
@@ -223,7 +210,7 @@ class Stegano_LSB_Decode:
                 temp_path = os.path.join(temp_dir, f"image_{i}.png")
                 img_pil.save(temp_path)
                 generator_type = generator_func(temp_path)
-            final_output.append(stegano.lsb.reveal(img_pil, generator_type, encoding=encoding))
+            final_output.append(stegano.lsb.reveal(img_pil, generator_type, encoding=encoding))  # ty: ignore[invalid-argument-type]
         return ("".join(final_output),)
 
 

src/utils.py

@@ -133,7 +133,7 @@ class Base642String(ConverterNodes):
         return (base64.b64decode(text.encode("ascii")).decode(self.encoding_selector(encoding_format, other_encoding_format)),)
 
 
-# The Bitwise category of Utilities. Also has auto-detection for the sake of my sanity.
+# The Bitwise category of Utilities.
 
 
 class BitwiseNodes:
@@ -161,7 +161,7 @@ class BitwiseNodes:
                         "default": "Hello World!",
                         "multiline": True,
                         "placeholder": "Type your message here...",
-                        "tooltip": "The secondary string to compare against. Accepts format defined in `datatype`",
+                        "tooltip": "The secondary string to compare against. Accepts format defined in `datatype`. Will be ignored in bitwise NOT operations.",
                     },
                 ),
                 "datatype": (["String", "Hexadecimal", "Base64", "Integer", "Binary"], {"default": "Binary", "tooltip": "The datatype that the two texts is assumed to be for comparison purpose. Will also output the final output based on the chosen datatype as well."}),
@@ -202,7 +202,7 @@ class BitwiseNodes:
             encoding_format = other_encoding_format
         return encoding_format
 
-    def detect_and_parse(self, text, datatype, encoding_format, other_encoding_format=None):
+    def detect_and_parse(self, text: str, datatype, encoding_format, other_encoding_format=None):
         text = text.strip()
 
         if datatype == "Binary":
@@ -231,12 +231,12 @@ class BitwiseNodes:
         else:
             raise ValueError(f"Unsupported datatype: {datatype}")
 
-    def operate(self, text_1, text_2, datatype, encoding_format, other_encoding_format=None, func=None):
-        b1, return_fmt = self.detect_and_parse(
+    def operate(self, text_1: str, text_2: str, datatype, encoding_format, func, other_encoding_format=None):
+        b1, return_fmt= self.detect_and_parse(
             text=text_1, datatype=datatype, encoding_format=encoding_format, other_encoding_format=other_encoding_format
         )
 
-        if text_2 is not None:
+        if text_2 is not "":
             b2, _ = self.detect_and_parse(
                 text=text_2, datatype=datatype, encoding_format=encoding_format, other_encoding_format=other_encoding_format
             )
@@ -299,7 +299,7 @@ class BitwiseNOT(BitwiseNodes):
     def execute(self, text, encoding_format, other_encoding_format):
         result_bytes, input_format = self.operate(
             text,
-            None,
+            "",
             encoding_format,
             other_encoding_format,
             func=lambda b1: bytes((~x & 0xFF) for x in b1),
@@ -374,8 +374,8 @@ class StringLooper:
     RETURN_NAMES = ("looped_text",)
     FUNCTION = "Looper"
 
-    def Looper(self, text=str, loops=int):
-        return (text * (loops + 1),)
+    def Looper(self, text: str, loops: int):
+        return (text*loops,)
 
 
 class ByteslikeEncode:
@@ -419,7 +419,7 @@ class ByteslikeEncode:
                         raise ValueError("Input for 'Raw Bytes' must evaluate to a bytes object (e.g., b'text').")
                 case _:
                     raise ValueError(f"Invalid encoding method used: {encoding}.")
-        except [ValueError, TypeError, UnicodeDecodeError] as err:
+        except (ValueError, TypeError, UnicodeDecodeError) as err:
             logger.warning("Bytes-like Object Encode at node ID %s has failed to encode the string given under the mode of \"%s\" with %s. Falling back to raw string bytes.", unique_id, encoding, err)
             data = text.encode("utf-8")
         return (data,)
@@ -430,7 +430,7 @@ class ByteslikeDecode:
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "data": ("BYTESLIKE",{}), 
+                "data": ("BYTESLIKE",{}),
                 "encoding": (["Hexadecimal", "Base64", "UTF-8", "Binary", "Raw Bytes"], {"default": "UTF-8", "tooltip": "The type of data the bytes will be decoded to."})
                 },
             "optional": {},