4
\$\begingroup\$

I am new to Python. Amateur, enthusiast and learner.
I have developed code which visits one particular website. Crawls through it to reach a certain repository, downloads the records from there to the local disk. It has too many while- and for-loops nested inside one another. It has nested functions too. How can I refactor it to increase speed and readability? Pasting the code here in hope of guidance.

import glob
import datetime
import cv2
import base64
from PIL import Image
from io import BytesIO
import time
import selenium
import self as self
from pytesseract import pytesseract
from selenium.webdriver.common.keys import Keys
import os
from selenium.webdriver.support import expected_conditions as EC, expected_conditions
from selenium.webdriver.support.select import Select
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.common.by import By
from selenium.common.exceptions import NoSuchElementException, TimeoutException, StaleElementReferenceException, \
    WebDriverException, ElementNotInteractableException, UnexpectedAlertPresentException

main_Directory = r'/home/sangharshmanuski/Documents/e_courts/mha/downloads4'
log_Directory = r'/home/sangharshmanuski/Documents/e_courts/mha/log'
driver = selenium.webdriver.Firefox()
url = r'https://districts.ecourts.gov.in/'
driver.get(url)
# create wait time variable for regular, short and mid
wait = WebDriverWait(driver, 180)
waitShort = WebDriverWait(driver, 3)
wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, "#sateist > option:nth-child(22)")))
select = Select(driver.find_element_by_css_selector('#sateist'))
options = select.options
select.select_by_visible_text('Maharashtra')
wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, '.region')))
wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, '#sateist')))
districtListDropdown = Select(driver.find_element_by_css_selector("#sateist"))
distOptions = districtListDropdown.options

# iterate over each district
i = 1
while i < len(distOptions):
    try:
        wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, '#sateist')))
        newDistDropDown = Select(driver.find_element_by_css_selector("#sateist"))
    except:
        continue
    newDistOptions = newDistDropDown.options
    distName = newDistOptions[i].text
    print(distName)
    newDistDropDown.select_by_index(i)
    # for creating directory as per each district.
    district_directory = os.path.join(
        main_Directory, distName)  # create new
    if not os.path.exists(district_directory):  # if not directory exists, create one
        os.mkdir(district_directory)
    district_log_directory = os.path.join(log_Directory, distName)
    if not os.path.exists(district_log_directory):  # if not directory exists, create one
        os.mkdir(district_log_directory)
    headingDist = driver.find_element_by_css_selector('.heading')
    if headingDist.text.lower() == distName.lower():
        wait.until(
            EC.element_to_be_clickable((By.CSS_SELECTOR, 'button.accordion2:nth-child(2)'))).click()
        current = driver.window_handles[0]
        wait.until(EC.element_to_be_clickable(
            (By.CSS_SELECTOR,
             'div.panel:nth-child(3) > ul:nth-child(1) > li:nth-child(6) > a:nth-child(1)'))).click()
        # wait until new tab opens.
        wait.until(EC.number_of_windows_to_be(2))
        # define new tab by differentiating from current tab.
        newWindow = [window for window in driver.window_handles if window != current][0]
        # switch to the new tab. ref: https://stackoverflow.com/questions/41571217/python-3-5-selenium-how-to-handle-a-new-window-and-wait-until-it-is-fully-lo
        driver.switch_to.window(newWindow)
        # wait till court complex list appears.
        wait.until(EC.element_to_be_clickable((By.CSS_SELECTOR, '#court_complex_code')))
        # create list of all court complex.
        # 2 approaches - 1 select 2 click.
        time.sleep(3)


        def complex_and_act():
            this = driver.current_window_handle

            def imgtotxt():
                elem = driver.find_element_by_id("captcha_image")
                loc = elem.location
                size = elem.size
                left = loc['x']
                top = loc['y']
                width = size['width']
                height = size['height']
                box = (int(left), int(top), int(left + width), int(top + height))
                screenshot = driver.get_screenshot_as_base64()
                img = Image.open(BytesIO(base64.b64decode(screenshot)))
                area = img.crop(box)
                area.save('/home/sangharshmanuski/Documents/e_courts/captcha/file_trial.png', 'PNG')
                fullPath = r'/home/sangharshmanuski/Documents/e_courts/captcha'
                f = os.listdir(fullPath)
                desPath = r"/home/sangharshmanuski/Documents/e_courts/editC"
                img = cv2.imread(os.path.join(fullPath, 'file_trial.png'), 0)
                ret, thresh1 = cv2.threshold(img, 111, 255, cv2.THRESH_BINARY)
                cv2.imwrite('/home/sangharshmanuski/Documents/e_courts/editC/oneDisNoLoop.png', thresh1)
                # know the text with pytesseract
                captchaText = pytesseract.image_to_string(
                    Image.open('/home/sangharshmanuski/Documents/e_courts/editC/oneDisNoLoop.png'))
                captcha = driver.find_element_by_id('captcha')
                captcha.send_keys(captchaText)
                driver.find_element_by_css_selector('input.button:nth-child(1)').click()
                time.sleep(1)

            def proceed():
                while True:
                    try:
                        waitShort.until(EC.alert_is_present())
                        driver.switch_to.alert.accept()
                        driver.switch_to.window(this)
                        driver.find_element_by_css_selector(
                            '#captcha_container_2 > div:nth-child('
                            '1) > div:nth-child(1) > span:nth-child(3) > a:nth-child(7) > img:nth-child(1)').click()
                        log_file = open(os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                        log_file.write('alert was present' + '\n')
                        print('alert was present')
                        imgtotxt()
                    except:
                        # if the waitmsg is on, wait for 5 sec
                        log_file = open(os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                        log_file.write('no alert' + '\n')
                        print('no alert')
                        waitmsg = 0
                        while driver.find_element_by_css_selector('#waitmsg').is_displayed():
                            if waitmsg < 7:
                                log_file = open(
                                    os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                                log_file.write('wait' + '\n')
                                print('waitmsg')
                                time.sleep(1)
                                waitmsg += 1
                            else:
                                log_file = open(os.path.join(
                                    log_Directory, nameCourtComp + '.txt'), 'a')
                                log_file.write('waiting finished' + '\n')
                                print('waiting finished')
                                break
                        invalidCaptcha = "Invalid Captcha"
                        norecord = "Record Not Found"
                        try:
                            waitShort.until(
                                EC.presence_of_element_located((By.CSS_SELECTOR, '#errSpan > p:nth-child(1)')))
                            incorrect = driver.find_element_by_css_selector('#errSpan > p:nth-child(1)').text
                            if incorrect == invalidCaptcha:
                                log_file = open(
                                    os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                                log_file.write('Invalid Captcha' + '\n')
                                print('invalid captcha')
                                imgtotxt()
                                continue
                            else:
                                if incorrect == norecord:
                                    log_file = open(
                                        os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                                    log_file.write('Record not Found' + '\n')
                                    return print('record not found')

                        except:
                            pass

                        def record():
                            log_file = open(
                                os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                            log_file.write('Record Found' + '\n')
                            print('record fun started')
                            wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, 'a.someclass')))
                            listAllView = driver.find_elements_by_css_selector(
                                'a.someclass')
                            # make new dirctory by name of Court Complex
                            distDir2 = os.path.join(
                                main_Directory, distName, nameCourtComp)
                            if not os.path.exists(distDir2):
                                os.makedirs(distDir2)
                            x = 0
                            for view in listAllView:
                                try:
                                    view.click()
                                    wait.until(EC.presence_of_element_located((By.ID, 'back_top')))
                                    openFile = open(
                                        os.path.join(distDir2, "file_" + str(x) + ".html"), "w")
                                    openFile.write(driver.page_source)
                                    openFile.close()
                                    back = driver.find_element_by_id('back_top')
                                    back.click()
                                    x += 1
                                except (TimeoutException, ElementNotInteractableException):

                                    driver.refresh()
                                    log_file = open(
                                        os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                                    log_file.write(
                                        'While Downloading record for '
                                        + nameCourtComp + ' error occured, retrying now...' + '\n')
                                    nonlocal courtComp
                                    courtComp -= 1
                                    return print(
                                        'While Downloading record for '
                                        + nameCourtComp + ' error occured, retrying now...')
                            log_file = open(
                                os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                            log_file.write('record completed, ' + str(x) + ' records found' + '\n')
                            print('record completed, ' + str(x) + ' records found')
                            return

                        record()
                        return
            courtComp = 1
            courtComplexDownload = Select(
                driver.find_element_by_css_selector('#court_complex_code'))
            courtComplexDownloadList = courtComplexDownload.options
            courtComplexLen = len(courtComplexDownloadList)
            while courtComp < courtComplexLen:
                nameCourtComp = courtComplexDownloadList[courtComp].text
                log_file = open(os.path.join(log_Directory, nameCourtComp + '.txt'), 'w')
                log_file.write(nameCourtComp + '\n' + '\n')
                print(nameCourtComp)
                courtComplexDownload.select_by_index(courtComp)
                acts = Select(driver.find_element_by_css_selector('#actcode'))
                actsOpt = acts.options
                act = 0
                while len(actsOpt) < 2:
                    if act < 10:

                        time.sleep(1)
                        act += 1
                    else:
                        #if there is no list to populate break out of this loop & go to next complex
                        raise Exception()
                try:
                    acts.select_by_value('33')
                except NoSuchElementException:
                    print('PoA not applicable')
                    log_file = open(
                        os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')
                    log_file.write('No PoA' + '\n')
                    courtComp += 1
                    continue

                imgtotxt()
                proceed()
                courtComp += 1


        complex_and_act()
        driver.close()
        print("all court complexes in " + distName + " completed")
        driver.switch_to.window(current)
        driver.back()

    else:
        time.sleep(5)
        continue
    i += 1

wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, "#sateist > option:nth-child(22)")))
select = Select(driver.find_element_by_css_selector('#sateist'))
options = select.options
select.select_by_visible_text('Maharashtra')
wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, '.region')))
wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, '#sateist')))
\$\endgroup\$
  • 1
    \$\begingroup\$ Technically, if it becomes faster, that's more than just refactoring. With that said, any task this thorny will probably be accomplished by rotating between refactoring & other tasks it helps you work out. \$\endgroup\$ – J.G. May 8 at 20:00
  • \$\begingroup\$ @J.G. That's not right. Refactoring can lead to speed increases. Granted that changes specifically aimed at increasing speed will probably involve more than refactoring. \$\endgroup\$ – JBentley May 9 at 12:32
  • \$\begingroup\$ @JBentley I suppose it comes down to whether you consider speed part of "behaviour". Personally, I'd prefer refactoring to mean "primarily focus on non-functional improvements, but with maybe a bit of functional improvement along the way", not that I'm expecting "refactoring-oriented programming" to ever catch on as a paradigm. ;) \$\endgroup\$ – J.G. May 9 at 12:46
  • \$\begingroup\$ @J.G. Yes I agree, but my point is that changes that don't modify behaviour (i.e. pure refactoring changes) can lead to speed changes. That's because speed is a result of compiler/interpreter output, not a result of behaviour. Compiler output won't necessarily be identical for two pieces of behaviourally equivalent code. \$\endgroup\$ – JBentley May 9 at 15:05
  • \$\begingroup\$ @JBentley Ah, I thought you meant more radical algorithm tweaks that get people to break out big O notation. \$\endgroup\$ – J.G. May 9 at 15:07
4
\$\begingroup\$

So a few tips:

Generally, functions go at the top level or as object or class methods. But doing that does mean you have to pass more things to the function, and can't rely on the closure (variables defined in the scope above the function). That can also give you a lot of insight into your structure, and highlight complex functions (complex means dealing with a lot of different things) you might want to simplify, so it's a good exercise to move those functions to the top. If you find a lot of functions need the same variable(s), then you might want a class. Or maybe you do decide you want the function nested.

Surprisingly though, I think you have too few functions, not too many. Some may be in the wrong places though. Most of the places you have a comment are also good places to extract a function. Don't be afraid to do that too much. More functions are usually better, and shorter functions that call other functions are great too. In fact, sometimes I extract a single line to a function just to give it a name, instead of using a comment. PyCharm makes it easy to (un)extract and move functions, so you can go hog wild pretty quickly. Anything that you can give a good function name is probably a good function.

Those suggestions lead to a pattern though: Page Objects. Basically, you make classes and objects to represent each conceptual element (things like header, logo, article content, login form, etc., not things like div, span, etc.) on the page, then use those objects to as your high level interface. So for example you could use that interface like this:

page = Page(url)

page.header.login_form.should_exist

page.main_content.should_contain("Welcome!")

next_page = page.header.login_form.fill_out_and_submit(name="J. Doe", password=password)

next_page.main_content.should_contain("Welcome J. Doe!")

The general idea is you want most of your code to read like a story. If you can simplify the story by using a function name that clearly describes a few lines, then that is probably a good idea. The details of how selenium works should be hidden from that story.

| improve this answer | |
\$\endgroup\$
2
\$\begingroup\$

You'll have to test all this, but:

  • Your editor highlighting will guide you. When I pulled out your functions, I found a need to add this and courtComp as arguments to proceed and courtComp to record, but at least I could retire a nonlocal. I also found, going through your code, that some variables can be retired altogether.
  • You can define a few more functions to DRY your code, e.g.

def ensure_dir(dir): if not os.path.exists(dir): os.mkdir(dir)

and

def wait_(): select = Select(driver.find_element_by_css_selector('#sateist')) select.select_by_visible_text('Maharashtra') for s in ['.region', '#sateist']: wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, s)))

This also lets you retire some variable declarations.

  • To flatten (denest) your code further, put the shorter of two if/else options first, then the rest may not need the else. For example, if headingDist.text.lower() == distName.lower(): enters lots of code even after we pull out functions defined in it, but is followed by an else: that only enters two lines. The large number in that first block can each be unindented one level if you follow my advice.
  • If a block's more than a few lines, it's worth turning that into a function too for further denesting.
  • You should probably learn how to make code more Pythonic. For example, instead of saying i=1 and doing a while loop ending in i += 1, use a for loop with range. (It doesn't make a big difference in a long block, but that's kind of the point: the blocks will quickly become short with the above points.) You might also want to learn about f-strings.

None of this advice is intended to make it faster, but better understanding Selenium might.

| improve this answer | |
\$\endgroup\$
2
\$\begingroup\$

I will not offer a complete refactor of your script, it's too much work I'm afraid. But well done, this is not an easy project.

I also agree that you should have more functions. Break up your code in smaller bits, to separate functionality and make it more readable. The flow is not very easy to follow.

I haven't tried all of your code but if I understand the desired functionality you should have a few logical blocks like:

  • open the main page
  • pick a state
  • pick a district
  • collect some data
  • then come back to the home page and repeat the process

When you have a clear operating flow in mind you can start writing dedicated functions for the various tasks you've identified. I will give an example that retrieves the list of States/UT.
NB: tested on Linux with Firefox.
Warning: I may have changed a few options/imports from your original code.

from typing import List
from selenium import webdriver
from selenium.common.exceptions import TimeoutException
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support.ui import Select
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from selenium.webdriver.firefox.options import Options as FirefoxOptions

# constants
url = r'https://districts.ecourts.gov.in/'

options = FirefoxOptions()
#options.add_argument("--headless")
options.add_argument("--private-window")
driver = webdriver.Firefox(options=options)

# FUNCTIONS
def get_states(driver) -> List[str]:
    """Get list of States/UT from combo box
    Return a list of strings
    """

    # define the selector only once
    combo_identifier = '#sateist'

    try:
        # wait for combo box to be ready
        print("Waiting for combo box (States/UT)...")
        WebDriverWait(driver=driver, timeout=10).until(EC.presence_of_element_located((By.CSS_SELECTOR, combo_identifier)))
        print("Combo box should be ready, continue")
    except TimeoutException:
        print("Timed out/failed to load page")
        sys.exit()

    states_combo = Select(driver.find_element_by_css_selector(combo_identifier))

    # return list of non-empty values from combo box
    return [o.get_attribute("value") for o in states_combo.options if o.get_attribute("value") != '']


# MAIN CODE

# load the main page
driver.get(url)

# Step 1 - get the list of States/UT
print("List of States/UT:")
for counter, value in enumerate(get_states(driver=driver), start=1):
    print(f'[{counter}] {value}')
    # Step 2: choose a district
    # ....

Output:

List of States/UT:
Waiting for combo box (States/UT)...
Combo box should be ready, continue
[1] andaman
[2] ap
[3] arunachal
[4] assam
[5] bihar
[6] chandigarh-district-court
[7] chhattisgarh
[8] dadra
[9] damandiu
[10] delhi
[11] goa
[12] gujarat
[13] haryana
[14] hp
[15] jk
[16] jharkhand
[17] karnataka
[18] kerala
[19] lakshadweep
[20] mp
[21] maharashtra
[22] manipur
[23] meghalaya
[24] mizoram
[25] nagaland
[26] odisha
[27] puducherry
[28] punjab
[29] rajasthan
[30] sikkim
[31] tn
[32] telangana
[33] tripura
[34] up
[35] uttarakhand
[36] wb

The function get_states returns a list of string, combo values only. If you are interested in the State names as well, then you could return a a list of dictionary pairs.

Now that you have fetched the list you can select one item at a time, for example to choose Goa in the list you do this:

combo_identifier = '#sateist'
state_option = Select(driver.find_element_by_css_selector(combo_identifier))
state_option.select_by_value('goa')

Regarding structure/style I would recommend to add some more line spacing, and more comments to better follow the logic. I bet that the reason why you have so little line spacing is that your code is already so long and takes time to scroll. That's why you need small functions (say, 10-30 lines each) instead of big blocks.


Something else that may be useful for you: logging. Python has an elaborate logging module, that you can use instead of writing to text files.

I use it extensively in my applications to write text to the console AND to file. Note that you can write to multiple destinations but with different formats.

You repeat paths and log file names too often eg:

os.path.join(log_Directory, nameCourtComp + '.txt'), 'a')

The logging module should help you simplify this. When you harness the power of this module you will get rid of those prints. This is more flexible because you can change the level of detail you want instead of commenting a lot of print statements.

Here is some sample code for demonstration purposes:

import logging
import sys

log_file = '/home/anonymous/test.log'

# logging - source: https://stackoverflow.com/questions/13733552/logger-configuration-to-log-to-file-and-print-to-stdout
# Change root logger level from WARNING (default) to NOTSET in order for all messages to be delegated.
logging.getLogger().setLevel(logging.NOTSET)

# Add stdout handler, with level INFO
console = logging.StreamHandler(sys.stdout)
console.setLevel(logging.INFO)
formater = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s', "%Y-%m-%d %H:%M:%S")
console.setFormatter(formater)
logging.getLogger().addHandler(console)

# Add file handler, with level DEBUG
handler = logging.FileHandler(log_file)
handler.setLevel(logging.DEBUG)
formater = logging.Formatter('%(asctime)s\t%(filename)s\t%(lineno)s\t%(name)s\t%(funcName)s\t%(levelname)s\t%(message)s', "%Y-%m-%d %H:%M:%S")
handler.setFormatter(formater)
logging.getLogger().addHandler(handler)

logger = logging.getLogger(__name__)

# this line will appear on console and in the log file
logger.info("Application started")

# this line will only appear in the log file because level  = debug
logger.debug("Log some gory details here")

Once again, well done. But if think your code is too long. Maintenance and debugging is not going to be easy. You really need to break it up in a number of steps and call them in a logical order that is easy to follow.

The site will change at some point and you will have to review your code sooner or later. If you split the code in functions and add a lot of traces it will be easier to figure out (function name, line number) where the problem occurs.

Something that you could/should do is put all the functions in a separate module file, then import it as a module from your main code. This will declutter your application a lot.

| improve this answer | |
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.