Creating nested list comprehension of files starting with specific string

Question

I have a directory with 'sets' of files that start with a state name followed by 4 or 5 digits (typically indicating year). Each 'file set' contains 3 files a .txt, a .png, and a .jpg.

Example of files in directory:

California1998_reform_paper.txt 
California1998_reform_paper.pdf
California1998_reform_paper.jpg
California2002_waterdensity_paper.txt
California2002_waterdensity_paper.pdf
California2002_waterdensity_paper.jpg

Based on a users input I am trying to write some code that can put each of these file sets into a list of lists. Ultimately I would like to iterate over the list of lists. That said, I am not married to any one data type if a dictionary or something else may be more efficient.

I would like the user to be able to enter either:

• The state name i.e. 'California' - to get all files from California
  OR
  
• The state name + year i.e. 'California1998' to get all files from California 1998

import os
import regex

directory = #path to directory
input = 'California1998'

# Does input match proper format? If not error.
mm = regex.match('^([a-z]+)([0-9]{4,5})|^([a-z]+)', input)
dir = str(os.listdir(directory))
if mm.group(1):
    state = mm.group(1)
    number = mm.group(2)
    state_num = state + number
    fileset = regex.findall(state_num, dir)
elif mm.group(3):
    state = mm.group(3)
    fileset = regex.findall(state + r'[0-9]{4,5}', dir)
else:
    print('Put some error message here')

# Does input exist? If not error.
if len(fileset) > 0:
    fileset = tuple(set(sorted(fileset)))
else:
    print('Put some error message here')

# Get list of lists
state_num_files = [[file.path
    for file in os.scandir(directory)
    if file.name.startswith(state_num)]
    for state_num in fileset]
return state_num_files

Above is the code I have thus far. It first uses regex.match to check the input, then regex.findall to find all matching state + years. I then create a sorted() set() from this list, which is converted into a tuple() called fileset. The last bit of code is a nested list comprehension that produces the list of lists by iterating through all files in the directory and iterating through all the state + years in fileset.

It certainly works, but seems repetitive and slower than it needs to be. My goal is to increase efficiency and remove any unnecessary iteration.

Thoughts on improvements:

• Possibly replace each regex.findall with a nested list comprehension? and thus remove the state_num_files nested comprehension at the end of script?

Any thoughts are greatly appreciated!

Answer 1

Review

1. Bug on Capitalization

   mm = regex.match('^([a-z]+)([0-9]{4,5})|^([a-z]+)', input)
   

   This does not work for the given use case of California1998

   But it can be easily fixed by adjusting the regex to include [A-Za-z] capital letters

2. Stop overshadowing!

   You use multiple built-in keywords as variable names ie, input dir this makes it that the overshadowed functions can no longer be used further in the program

3. import regex?

   I think this should be import re, since that's what the library is called

4. Unnecesarry operations

   The sort and consecutive conversion do nothing at all

   Secondly you loop over the directory twice! This can be avoided by doing it one loop.

5. Don't print errors raise them

   And make sure the user get's back useful information in the error message this generic message is not very helpful :)

Alternative

Look into the glob module this does what you require,

>>> import glob
>>> glob.glob("./California1998*")
['./California1998_reform_paper.jpg', './California1998_reform_paper.txt', './California1998_reform_paper.pdf']

Answer 2

Welcome to Stack Overflow! Based on your code and what your were trying to do (and one suggestion from Ludisposed) here's a attempt where you only scan once the directory where your files are situated:

import re
import pathlib
import os
import collections

directory_of_source_file = pathlib.Path("Input files")
file_set_from_input = 'California1998'

state_name_matcher = re.compile(r'([a-z]+)(\d{4,5})?', re.IGNORECASE)
if state_name_found := state_name_matcher.match(file_set_from_input):
    state_wanted = state_name_found[1]
    set_number = state_name_found[2]
    set_number = '' if set_number is None else set_number

    files_found = directory_of_source_file.glob(f"{state_wanted}{set_number}*")
    files_found_strings = list(map(str, files_found))
    if set_number:
        files_found_formatted = [files_found_strings, ]
    else:
        strict_state_name_matcher = re.compile(rf'{os.sep}{state_wanted}(\d{{4,5}})', re.IGNORECASE)
        documents_collector = collections.defaultdict(list)
        for current_file in files_found_strings:
            if matching_document_found := strict_state_name_matcher.search(current_file):
                current_set_number = matching_document_found[1]
                documents_collector[current_set_number].append(current_file)
        files_found_formatted = list(documents_collector.values())
    for current_file_set in files_found_formatted:
        current_file_set.sort()
    print(files_found_formatted)
else:
    print('Put some error message here')

I hope I understood correctly what you were trying to do