From f4e2505c072f685d380a29d0389efef425ccc546 Mon Sep 17 00:00:00 2001 From: Josselin Date: Wed, 2 Oct 2019 14:07:07 +0200 Subject: [PATCH] slither-flat: Add support for --contract --- slither/tools/flattening/flattening.py | 31 +++++++++++++++++--------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/slither/tools/flattening/flattening.py b/slither/tools/flattening/flattening.py index 2f6a2ecd5..6acd8200a 100644 --- a/slither/tools/flattening/flattening.py +++ b/slither/tools/flattening/flattening.py @@ -125,22 +125,31 @@ class Flattening: exported.add(contract.name) list_contract.append(self._source_codes[contract]) + def _export(self, contract, ret): + self._export_contract(contract, set(), ret) + path = Path(self.DEFAULT_EXPORT_PATH, f'{contract.name}.sol') + logger.info(f'Export {path}') + with open(path, 'w') as f: + if self._slither.solc_version: + f.write(f'pragma solidity {self._slither.solc_version};\n') + if self._use_abi_encoder_v2: + f.write('pragma experimental ABIEncoderV2;\n') + f.write('\n'.join(ret)) + f.write('\n') def export(self, target=None): if not self.DEFAULT_EXPORT_PATH.exists(): self.DEFAULT_EXPORT_PATH.mkdir(parents=True) + ret = [] if target is None: for contract in self._slither.contracts_derived: - ret = [] - self._export_contract(contract, set(), ret) - path = Path(self.DEFAULT_EXPORT_PATH, f'{contract.name}.sol') - logger.info(f'Export {path}') - with open(path, 'w') as f: - if self._slither.solc_version: - f.write(f'pragma solidity {self._slither.solc_version};\n') - if self._use_abi_encoder_v2: - f.write('pragma experimental ABIEncoderV2;\n') - f.write('\n'.join(ret)) - f.write('\n') + self._export(contract, ret) + else: + contract = self._slither.get_contract_from_name(target) + if contract is None: + logger.error(f'{target} not found') + else: + self._export(contract, ret) +