From af4d7d87f39ff4b24a2d627ea93f186ef1c78e3f Mon Sep 17 00:00:00 2001
From: insolor <2442833+insolor@users.noreply.github.com>
Date: Tue, 5 Nov 2024 09:29:04 +0300
Subject: [PATCH] Refactoring

---
 dfint64_patch/extract_strings/cli.py                       | 3 +--
 dfint64_patch/strings_context/extract_strings_with_subs.py | 6 +++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/dfint64_patch/extract_strings/cli.py b/dfint64_patch/extract_strings/cli.py
index 4d111de..e4c5577 100644
--- a/dfint64_patch/extract_strings/cli.py
+++ b/dfint64_patch/extract_strings/cli.py
@@ -1,5 +1,4 @@
 import operator
-from collections.abc import Iterable
 from dataclasses import dataclass
 from pathlib import Path
 from typing import BinaryIO, cast
@@ -21,7 +20,7 @@
 from dfint64_patch.utils import maybe_open
 
 
-def extract_strings(pe_file: BinaryIO) -> Iterable[ExtractedStringInfo]:
+def extract_strings(pe_file: BinaryIO) -> list[ExtractedStringInfo]:
     pe = PortableExecutable(pe_file)
 
     sections = pe.section_table
diff --git a/dfint64_patch/strings_context/extract_strings_with_subs.py b/dfint64_patch/strings_context/extract_strings_with_subs.py
index ca8c8e9..785753c 100644
--- a/dfint64_patch/strings_context/extract_strings_with_subs.py
+++ b/dfint64_patch/strings_context/extract_strings_with_subs.py
@@ -1,5 +1,5 @@
 """
-TODO: Extract strings grouped by subroutines
+Extract strings grouped by subroutines
 """
 from collections import defaultdict
 from dataclasses import dataclass
@@ -48,7 +48,7 @@ class StringCrossReference(NamedTuple):
     cross_reference: Rva
 
 
-def extract_strings_with_subs(pe_file: BinaryIO) -> dict[Rva, list[StringCrossReference]]:
+def extract_strings_grouped_by_subs(pe_file: BinaryIO) -> dict[Rva, list[StringCrossReference]]:
     pe = PortableExecutable(pe_file)
     sections = pe.section_table
     code_section = sections[0]
@@ -87,7 +87,7 @@ class ExtractConfig(DictConfig):
 @with_config(ExtractConfig, ".extract.yaml")
 def main(conf: ExtractConfig) -> None:
     with Path(conf.file_name).open("rb") as pe_file:
-        for subroutine, strings in extract_strings_with_subs(pe_file).items():
+        for subroutine, strings in extract_strings_grouped_by_subs(pe_file).items():
             print(f"sub_{subroutine:x}:")
             for string in strings:
                 print(f"\t{string.string}")