Making custom class to behave like set

问题:

I try to make a class containing file names in a folder. But I want it to behave like set. Now I have this:

class Files():

    def __init__(self, in_dir):
        self.in_dir = in_dir
        self.files = set(map(os.path.basename, glob.glob(self.in_dir + "/*.txt")))

    def __add__(self, other):
        return self.files + other.files    

    def __or__(self, other):
        return self.files | other.files

    def __and__(self, other):
        return self.files & other.files

    def __xor__(self, other):
        return self.files ^ other.files

This work and I can do like this:

f1 = Files(inDir1)
f2 = Files(inDir2)

diff_files = f1 ^ f2 % this give files that are in f1 or f2 folder but not in both  folders

This is ok, but the problem is that diff_files is not instance of Files. How to change my class, to behave like set in python 3.x?

回答1:

First, make in_dir argument optional:

def __init__(self, in_dir=None):
    if in_dir:
        self.in_dir = in_dir
        self.files = set(map(os.path.basename, glob.glob(self.in_dir + "/*.txt")))

Then, change the __xor__():

def __xor__(self, other):
    instance = Files()
    instance.files = self.files ^ other.files
    return instance

Also, I don't see the reason to keep in_dir as an instance variable. You can simplify the __init__():

def __init__(self, in_dir=None):
    if in_dir:
        self.files = set(map(os.path.basename, glob.glob(in_dir + "/*.txt")))

---

Alternatively, you can allow to initialize Files by passing a files set:

def __init__(self, in_dir=None, files=None):
    if in_dir:
        self.files = set(map(os.path.basename, glob.glob(in_dir + "/*.txt")))
    if files:
        self.files = files

Then, the __xor__() method would be even simpler:

def __xor__(self, other):
    return Files(files=self.files ^ other.files)

回答2:

I'm not sure I understand what you mean by "behaves like set" but I do understand that you want to return an instance of Files and not only the "diff", so for that:

change:

def __xor__(self, other):
        return self.files ^ other.files

to:

def __xor__(self, other):
        result = Files()
        result.in_dir = self.in_dir
        result.files = self.files ^ other.files
        return result