Creating nested list comprehension of files starting with specific string 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) &gt; 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!
Review


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
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
import regex? 

I think this should be import re, since that's what the library is called
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.
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, 

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