Provide source content via crytic compile

crytic · Jun 20, 2019 · d9846f2 · d9846f2
1 parent 86b7cd6
commit d9846f2
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/crytic_compile/crytic_compile.py b/crytic_compile/crytic_compile.py
@@ -257,15 +257,13 @@ def bytecode_init(self, name, libraries=None):
         init = self._init_bytecodes.get(name, None)
         return self._update_bytecode_with_libraries(init, libraries)
 
-
     # endregion
     ###################################################################################
     ###################################################################################
     # region Source mapping
     ###################################################################################
     ###################################################################################
 
-
     @property
     def srcmaps_init(self):
         return self._srcmaps
@@ -291,6 +289,9 @@ def src_content(self):
                         self._src_content[filename.absolute] = source_file.read()
         return self._src_content
 
+    def src_content_for_file(self, filename_absolute):
+        return self.src_content.get(filename_absolute, None)
+
     # endregion
     ###################################################################################
     ###################################################################################
@@ -544,7 +545,10 @@ def _import_archive_compilation(self, compiled_archive, compilation_index):
 
             archive.set_dependency_status(filename.absolute, contract['is_dependency'])
 
-        # Set all our filenames
+            # Set all source code
+            self._src_content[filename.absolute] = compiled_archive['source_content'][filename.absolute]
+
+        # Set our filenames
         self._filenames = set(self._contracts_filenames.values())
 
         self._working_dir = compilation['working_dir']
@@ -572,7 +576,7 @@ def export_all(compilations, **kwargs):
         # Determine if we are exporting a singular archive, or exporting each individually.
         if export_format == 'archive':
             # Create our source file dictionary
-            results['source_files'] = {}
+            results['source_content'] = {}
 
             # If we are to export source..
             for compilation in compilations:
@@ -581,7 +585,7 @@ def export_all(compilations, **kwargs):
 
                 # Next set all source content
                 for filename_absolute, source_content in compilation.src_content.items():
-                    results['source_files'][filename_absolute] = source_content
+                    results['source_content'][filename_absolute] = source_content
 
             # If we have an export directory specified, we output the JSON to a file.
             export_dir = kwargs.get('export_dir', None)