From 5287b0aebf3452b174c1910ff903f43b48aba683 Mon Sep 17 00:00:00 2001
From: Josselin Feist <josselin.feist@gmail.com>
Date: Thu, 4 Aug 2022 16:26:21 +0200
Subject: [PATCH] Minor

---
 slither/tools/read_storage/__main__.py     | 23 +++++++++++-----------
 slither/tools/read_storage/read_storage.py |  2 +-
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/slither/tools/read_storage/__main__.py b/slither/tools/read_storage/__main__.py
index 435a138723..f3cc90efca 100644
--- a/slither/tools/read_storage/__main__.py
+++ b/slither/tools/read_storage/__main__.py
@@ -23,7 +23,7 @@ def parse_args() -> argparse.Namespace:
             + "To retrieve a contract's storage layout:\n"
             + "\tslither-read-storage $TARGET address --contract-name $NAME --layout\n"
             + "To retrieve a contract's storage layout and values:\n"
-            + "\tslither-read-storage $TARGET address --contract-name $NAME --layout --values\n"
+            + "\tslither-read-storage $TARGET address --contract-name $NAME --layout --value\n"
             + "TARGET can be a contract address or project directory"
         ),
     )
@@ -73,9 +73,9 @@ def parse_args() -> argparse.Namespace:
     )
 
     parser.add_argument(
-        "--layout",
-        action="store_true",
-        help="Toggle used to write a JSON file with the entire storage layout.",
+        "--json",
+        action="store",
+        help="Save the result in a JSON file.",
     )
 
     parser.add_argument(
@@ -132,15 +132,14 @@ def main() -> None:
 
         srs.rpc = args.rpc_url
 
-    if args.layout:
-        srs.get_all_storage_variables()
-        srs.get_storage_layout()
-    else:
-        assert args.variable_name
+    if args.variable_name:
         # Use a lambda func to only return variables that have same name as target.
         # x is a tuple (`Contract`, `StateVariable`).
         srs.get_all_storage_variables(lambda x: bool(x[1].name == args.variable_name))
         srs.get_target_variables(**vars(args))
+    else:
+        srs.get_all_storage_variables()
+        srs.get_storage_layout()
 
     # To retrieve slot values an rpc url is required.
     if args.value:
@@ -150,9 +149,9 @@ def main() -> None:
     if args.table:
         srs.walk_slot_info(srs.convert_slot_info_to_rows)
         print(srs.table)
-    # Only write file if storage layout is used. TODO add flag for file
-    elif len(srs.slot_info) > 1:
-        with open("storage_layout.json", "w", encoding="utf-8") as file:
+
+    if args.json:
+        with open(args.json, "w", encoding="utf-8") as file:
             slot_infos_json = srs.to_json()
             json.dump(slot_infos_json, file, indent=4)
 
diff --git a/slither/tools/read_storage/read_storage.py b/slither/tools/read_storage/read_storage.py
index 78e0d610be..c9523f7330 100644
--- a/slither/tools/read_storage/read_storage.py
+++ b/slither/tools/read_storage/read_storage.py
@@ -106,7 +106,7 @@ def slot_info(self) -> Dict[str, SlotInfo]:
 
     def get_storage_layout(self) -> None:
         """Retrieves the storage layout of entire contract."""
-        tmp = {}
+        tmp: Dict[str, SlotInfo] = {}
         for contract, var in self.target_variables:
             type_ = var.type
             info = self.get_storage_slot(var, contract)