<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def write_output(self, data, file_name, key=None):
 """Writes output data to disk. Parameters data: pandas.DataFrame or dict The data to write to disk. file_name: str The name of the file to write. key: str, optional The lookup key for the sub_directory to write results to, if any. """

        path = os.path.join(self._directories[key], file_name)
        extension = file_name.split('.')[-1]

        if extension == 'yaml':
            with open(path, 'w') as f:
                yaml.dump(data, f)
        elif extension == 'hdf':
            # to_hdf breaks with categorical dtypes.
            categorical_columns = data.dtypes[data.dtypes == 'category'].index
            data.loc[:, categorical_columns] = data.loc[:, categorical_columns].astype('object')
            # Writing to an hdf over and over balloons the file size so write to new file and move it over to avoid
            data.to_hdf(path + "update", 'data')
            if os.path.exists(path):
                os.remove(path)
            os.rename(path + "update", path)
        else:
            raise NotImplementedError(
                f"Only 'yaml' and 'hdf' file types are supported. You requested {extension}")