I'm trying to get 4 files from a directory with an specific extension each one-CodePudding

files = [
    f for f in listdir(input_path)
    if path.isfile(path.join(input_path, f))
]
if files:
    for file in files:
        if file.endswith(".xml"):
            xml_filename = input_path   file
        elif file.endswith(".csv"):
            csv_filename = input_path   file
        elif file.endswith(".osgb"):
            osgb_filename = input_path   file
        elif file.endswith(".xodr"):
            xodr_filename = input_path   file

I'm trying to get 4 files from a directory with an specific extension each one but my solution looks kinda ugly you smart guys may have a clever solution ;D

CodePudding user response：

You can reduce code count if you move your result into a collection that can be filled in a loop. With individual variables, you need code per variable for the assignment. Using a dictionary and the standard pathlib module, your code could be

from pathlib import Path
files = {path.suffix[1:]:path for path in Path(input_path).iterdir()
    if path.suffix in {".csv", ".xml", ".osgb", ".xodr"}}

Now xml_filename is files["xml"].

CodePudding user response：

use glob.glob()

from glob import glob
import os

xml_filename = glob(os.path.join(input_path, '*.xml'))[0]
csv_filename = glob(os.path.join(input_path, '*.csv'))[0]
osgb_filename = glob(os.path.join(input_path, '*.osgb'))[0]
xodr_filename = glob(os.path.join(input_path, '*.xodr'))[0]

Note that this code assumes that there's at least one of each file type in the directory. You can use try/except to catch the IndexError if glob() doesn't return any matches.

CodePudding user response：

You could put your suffixes into a dict and compare them with the files suffix to get the filename:

suffs = {"xml": xml_filename, "csv": csv_filename, "osgb": osbg_filename, "xodr": xodr_filename}

if (fn := file.split('.')[-1]) in suff:
    suffs[fn] = input_path   file