From e96dbe7e5d162861ad77b279c5fec9b18db8cf78 Mon Sep 17 00:00:00 2001
From: Nate Harris <n8gr8gbln@gmail.com>
Date: Fri, 17 Apr 2020 20:52:41 -0400
Subject: [PATCH] Added CityStateLookup, JSON shortcut

---
 usps/address.py | 12 ++++++-----
 usps/usps.py    | 55 +++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/usps/address.py b/usps/address.py
index 7c84508..91bda2c 100644
--- a/usps/address.py
+++ b/usps/address.py
@@ -3,7 +3,7 @@
 
 class Address(object):
 
-    def __init__(self, name, address_1, city, state, zipcode,
+    def __init__(self, name, address_1, city, state, zipcode='',
                  zipcode_ext='', company='', address_2='', phone=''):
         self.name = name
         self.company = company
@@ -35,11 +35,13 @@ def add_to_xml(self, root, prefix='To', validate=False):
         state = etree.SubElement(root, prefix + 'State')
         state.text = self.state
 
-        zipcode = etree.SubElement(root, prefix + 'Zip5')
-        zipcode.text = self.zipcode
+        if self.zipcode:
+            zipcode = etree.SubElement(root, prefix + 'Zip5')
+            zipcode.text = self.zipcode
 
-        zipcode_ext = etree.SubElement(root, prefix + 'Zip4')
-        zipcode_ext.text = self.zipcode_ext
+        if self.zipcode_ext:
+            zipcode_ext = etree.SubElement(root, prefix + 'Zip4')
+            zipcode_ext.text = self.zipcode_ext
 
         if not validate:
             phone = etree.SubElement(root, prefix + 'Phone')
diff --git a/usps/usps.py b/usps/usps.py
index 63a4fc7..5ee954d 100644
--- a/usps/usps.py
+++ b/usps/usps.py
@@ -17,9 +17,11 @@ class USPSApi(object):
         'tracking': 'TrackV2{test}&XML={xml}',
         'label': 'eVS{test}&XML={xml}',
         'validate': 'Verify&XML={xml}',
+        'citystate': 'CityStateLookup&XML={xml}',
+        'zipcode': 'ZipCodeLookup&XML={xml}'
     }
 
-    def __init__(self,  api_user_id, test=False):
+    def __init__(self, api_user_id, test=False):
         self.api_user_id = api_user_id
         self.test = test
 
@@ -48,10 +50,16 @@ def track(self, *args, **kwargs):
     def create_label(self, *args, **kwargs):
         return ShippingLabel(self, *args, **kwargs)
 
+    def lookup_city_by_zip(self, *args, **kwargs):
+        return CityStateLookup(self, *args, **kwargs)
+
+    def lookup_zip_by_address(self, *args, **kwargs):
+        return ZipCodeLookup(self, *args, **kwargs)
+
 
 class AddressValidate(object):
 
-    def __init__(self, usps, address):
+    def __init__(self, usps, address, just_answer: bool = False):
         xml = etree.Element('AddressValidateRequest', {'USERID': usps.api_user_id})
         _address = etree.SubElement(xml, 'Address', {'ID': '0'})
         address.add_to_xml(_address, prefix='', validate=True)
@@ -68,6 +76,49 @@ def __init__(self, usps, tracking_number):
         self.result = usps.send_request('tracking', xml)
 
 
+class CityStateLookup(object):
+
+    def __init__(self, usps, zip_code, just_answer: bool = False):
+        xml = etree.Element('CityStateLookupRequest', {'USERID': usps.api_user_id})
+        _zip_code = etree.SubElement(xml, 'ZipCode', {'ID': '0'})
+        zip_code_inner = etree.SubElement(_zip_code, 'Zip5')
+        zip_code_inner.text = zip_code
+
+        self.result = usps.send_request('citystate', xml)
+        if just_answer:
+            try:
+                if self.result.get('CityStateLookupResponse').get('ZipCode').get('Error'):
+                    self.result = 'Invalid Zip Code.'
+                else:
+                    self.result = {
+                        'City': self.result.get('CityStateLookupResponse').get('ZipCode').get('City'),
+                        'State': self.result.get('CityStateLookupResponse').get('ZipCode').get('State')
+                    }
+            except:
+                self.result = None
+
+
+class ZipCodeLookup(object):
+
+    def __init__(self, usps, address, just_answer: bool = False):
+        xml = etree.Element('ZipCodeLookupRequest', {'USERID': usps.api_user_id})
+        _address = etree.SubElement(xml, 'Address', {'ID': '1'})
+        address.add_to_xml(_address, prefix='', validate=True)
+
+        self.result = usps.send_request('zipcode', xml)
+        if just_answer:
+            try:
+                if self.result.get('ZipCodeLookupResponse').get('Address').get('Error'):
+                    self.result = 'Invalid Address.'
+                else:
+                    self.result = {
+                        'Zip5': self.result.get('ZipCodeLookupResponse').get('Address').get('Zip5'),
+                        'Zip4': self.result.get('ZipCodeLookupResponse').get('Address').get('Zip4')
+                    }
+            except:
+                self.result = None
+
+
 class ShippingLabel(object):
 
     def __init__(self, usps, to_address, from_address, weight,