First of all, let's make use of the with
statement so that the file is closed automatically:
with open(path, 'r') as txt_file_object:
text = txt_file_object.readlines()
With this, you don't have to call close()
anymore, as the file will close automatically when you exit the with
scope.
text = text[0]
You are only reading the first line of text. Is this really what you want to do?
You are using the variable text
for two different things: for the input lines and for the output values. This is not very intuitive; in fact, the result can be a list of integers, so why would it be called text
? Maybe result
would be a better name for it.
BUT since now you don't have to close at the end of the function, you can return the result directly instead of saving it in a variable:
return [int(num) for num in text.split(sep)]
Returning will exit the with
scope, so again the file will be closed automatically.
[let for let in text.split(sep)]
This selects all objects inside text.split(sep)
, so we can return the splitted list directly:
text.split(sep)
Similarly, there's a different way of applying a function to every item in a list, which is using map
. Maybe the list comprehension feels more natural, so you can keep that if you want; still I'll show you just so you know:
# Creates a list by calling 'int' to every item in the list
list(map(int, text.split(sep))
You are repating a lot of code; you have four different results, but they are very similar, so let's try to provide a more generic way.
My concern here is handling the two possible dimensions. You are parsing the text in the same way (depending if it's int or str), but when it's two dimensions you do it for every line. So we could use Python's lambdas to decide first what type of parsing (int or str) we are doing, and then just apply it once or multiple times.
The lambda can take parameters; the only parameter we need in our case is the input text. We can't just use text
directly because sometimes we want to parse the full text, but sometimes only the line:
if typ == "int":
parse_function = lambda t: list(map(int, t.split(sep)))
elif typ == "str":
parse_function = lambda t: list(t.split(sep))
else:
raise ValueError("Unknown type.")
Now parse_function
can be used like any other function, taking the text as input. So we can use it when deciding the dimension:
if dim == 1:
return parse_function(text)
elif dim == 2:
return [parse_function(line) for line in text]
else:
raise ValueError("Unknown dimension.")
You do well in throwing exceptions for invalid input, but how is the user meant to know what possible values can be used for typ
and dim
? You could add that to the docstring. You should also say what the function does in the docstring.
Updated code
def read_array_from_txt(path, dim, typ, sep):
'''
Processes a text file as a 1D or 2D array.
:param path: Path to the input file.
:param dim: How many dimensions (1 or 2) the array has.
:param typ: Whether the elements are read as 'int' or 'str'
:param sep: The text that is used to separate between elements.
@time: O(n)
@space: O(n)
'''
with open(path,"r") as txt_file_object:
text = txt_file_object.readlines()
if typ == "int":
parse_function = lambda t: list(map(int, t.split(sep)))
elif typ == "str":
parse_function = lambda t: list(t.split(sep))
else:
raise ValueError("Unknown type.")
if dim == 1:
return parse_function(text)
elif dim == 2:
return [parse_function(line) for line in text]
else:
raise ValueError("Unknown dimension.")