Lint and typecheck.

Author: WhatTheFuzz

src/agent.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+from collections.abc import Callable
 import os
-from typing import Callable, Optional, Union
+from typing import Optional, Union
 from pathlib import Path
 
 import openai
@@ -67,7 +69,7 @@ def read_api_key(self, filename: Optional[Path]=None) -> str:
         settings: Settings = Settings()
         if settings.contains('openai.api_key'):
             if key := settings.get_string('openai.api_key'):
-                return key
+                return str(key)
 
         # If the settings don't exist, contain the key, or the key is empty,
         # check the environment variable.
@@ -106,7 +108,7 @@ def get_model(self) -> str:
             if model := settings.get_string('openai.model'):
                 # Check that is a valid model by querying the OpenAI API.
                 if self.is_valid_model(model):
-                    return model
+                    return str(model)
         # Return a valid, default model.
         assert self.is_valid_model('text-davinci-003')
         return 'text-davinci-003'
@@ -119,7 +121,7 @@ def get_token_count(self) -> int:
         if settings.contains('openai.max_tokens'):
             # Check that the value is not None.
             if (max_tokens := settings.get_integer('openai.max_tokens')) is not None:
-                return max_tokens
+                return int(max_tokens)
         return 1_024
 
     def instruction_list(self, function: Union[LowLevelILFunction,
@@ -161,7 +163,7 @@ def generate_query(self, function: Union[Function,
         return prompt
 
     def generate_rename_variable_query(self,
-                                    instruction: HighLevelILInstruction) -> Optional[str]:
+                                    instruction: HighLevelILInstruction) -> str:
         '''Generates a query string given a BNIL instruction. Returns the query
         as a string.
         '''

src/entry.py

@@ -1,26 +1,25 @@
 from pathlib import Path
 from binaryninja import BinaryView, Function
-from binaryninja.highlevelil import HighLevelILInstruction, HighLevelILVarInit, \
-                                    HighLevelILFunction
+from binaryninja.highlevelil import HighLevelILInstruction, HighLevelILVarInit
 from binaryninja.log import log_error
 from . agent import Agent
 
 API_KEY_PATH = Path.home() / Path('.openai/api_key.txt')
 
-def check_function(bv: BinaryView, func: Function) -> bool:
+def check_function(bv: BinaryView, func: Function) -> None:
     agent: Agent = Agent(
         bv=bv,
         path_to_api_key=API_KEY_PATH
     )
     query: str = agent.generate_query(func)
     agent.send_query(query)
 
-def rename_variable(bv: BinaryView, instruction: HighLevelILInstruction) -> bool:
+def rename_variable(bv: BinaryView, instruction: HighLevelILInstruction) -> None:
 
     if not isinstance(instruction, HighLevelILVarInit):
         log_error(f'Instruction must be of type HighLevelILVarInit, got type: ' \
                   f'{type(instruction)}')
-        return False
+        return
 
     agent: Agent = Agent(
         bv=bv,

src/query.py

@@ -1,12 +1,14 @@
+from __future__ import annotations
+from collections.abc import Callable
+from typing import Optional
 import openai
-from typing import Callable
 from binaryninja.plugin import BackgroundTaskThread
-from binaryninja.log import log_debug
+from binaryninja.log import log_debug, log_info
 
 class Query(BackgroundTaskThread):
 
     def __init__(self, query_string: str, model: str,
-                 max_token_count: int, callback_function: Callable=None) -> None:
+                 max_token_count: int, callback_function: Optional[Callable]=None) -> None:
         BackgroundTaskThread.__init__(self,
                                       initial_progress_text="",
                                       can_cancel=False)
@@ -20,16 +22,16 @@ def run(self) -> None:
 
         log_debug(f'Sending query: {self.query_string}')
 
-        response: str = openai.Completion.create(
+        response = openai.Completion.create(
             model=self.model,
             prompt=self.query_string,
             max_tokens=self.max_token_count,
         )
         # Get the response text.
-        response: str = response.choices[0].text
+        result: str = response.choices[0].text
         # If there is a callback, do something with it.
         if self.callback:
-            self.callback(response)
+            self.callback(result)
         # Otherwise, assume we just want to log it.
         else:
-            log_info(response)
+            log_info(result)