From 33aa51c19a8127b41f171d1b02acb8f6d2995c33 Mon Sep 17 00:00:00 2001
From: tuxxy <john@nucypher.com>
Date: Thu, 12 Apr 2018 17:38:35 -0600
Subject: [PATCH] Add get_contract_data method to find contracts on the
 Registrar chain via address or name

Fix syntax error
---
 nkms/blockchain/eth/interfaces.py | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/nkms/blockchain/eth/interfaces.py b/nkms/blockchain/eth/interfaces.py
index c61299eb0..6ae9704c8 100644
--- a/nkms/blockchain/eth/interfaces.py
+++ b/nkms/blockchain/eth/interfaces.py
@@ -48,7 +48,7 @@ class Registrar:
         self.__registrar_filepath = registrar_filepath or self.__DEFAULT_REGISTRAR_FILEPATH
 
     @classmethod
-    def get_chains(cls, registrar_filepath: str=None):
+    def get_chains(cls, registrar_filepath: str=None) -> dict:
         """
         Returns a dict of Registrar objects where the key is the chain name and
         the value is the Registrar object for that chain.
@@ -57,7 +57,7 @@ class Registrar:
         filepath = registrar_filepath or self.__DEFAULT_REGISTRAR_FILEPATH
         instance = cls(registrar_filepath=filepath)
 
-        registrar_data = instance._read_registrar_file(filepath)
+        registrar_data = _read_registrar_file(filepath)
         chain_names = registrar_data.keys()
 
         chains = dict()
@@ -66,7 +66,7 @@ class Registrar:
                                      registrar_filepath=filepath)
         return chains
 
-    def enroll(self, contract_name: str, contract_address: str, contract_abi: list):
+    def enroll(self, contract_name: str, contract_address: str, contract_abi: list) -> None:
         """
         Enrolls a contract to the chain registrar by writing the abi information
         to the filesystem as JSON. This can also be used to update the info
@@ -89,9 +89,11 @@ class Registrar:
 
         __write_registrar_file(registrar_data, self.__registrar_filepath)
 
-    def get_chain_data(self):
+    def get_chain_data(self) -> dict:
         """
         Returns all data from the current registrar chain as a dict.
+        If no data exists for the current registrar chain, then it will raise
+        KeyError.
         If you haven't specified the chain name, it's probably the tester chain.
         """
         registrar_data = __read_registrar_file(self.__registrar_filepath)
@@ -100,3 +102,22 @@ class Registrar:
         except KeyError:
             raise KeyError("Data does not exist for chain '{}'".format(self._chain_name))
         return chain_data
+
+    def get_contract_data(self, identifier: str=None) -> dict:
+        """
+        Returns contract data on the chain as a dict given an `identifier`.
+        It first attempts to use identifier as a contract name. If no name is
+        found, it will attempt to use identifier as an address.
+        If no contract is found, it will raise NoKnownContract.
+        """
+        chain_data = self.get_chain_data()
+        if identifier in chain_data:
+            contract_data = chain_data[identifier]
+            return contract_data
+        else:
+            for contract_name, contract_data in chain_data.items():
+                if contract_data['addr'] == identifier:
+                    return contract_data
+        raise self.NoKnownContract(
+            "Could not identify a contract name or address with {}".format(identifier)
+        )