treelite
filesystem.h
1 
7 #ifndef TREELITE_COMMON_FILESYSTEM_H_
8 #define TREELITE_COMMON_FILESYSTEM_H_
9 
10 #include <dmlc/logging.h>
11 #include <treelite/common.h>
12 #include <vector>
13 #include <string>
14 #include <regex>
15 #include <cstdlib>
16 
17 #ifdef _WIN32
18 #define NOMINMAX
19 #include <windows.h>
20 #else
21 #include <unistd.h>
22 #include <errno.h>
23 #include <fcntl.h>
24 #include <sys/stat.h>
25 #include <sys/types.h>
26 #include <libgen.h>
27 #include <cstring>
28 #endif
29 
30 namespace treelite {
31 namespace common {
32 namespace filesystem {
33 
42 inline std::string GetBasename(const std::string& path) {
43 #ifdef _WIN32
44  /* remove any trailing backward or forward slashes
45  (UNIX does this automatically) */
46  std::string path_;
47  std::string::size_type tmp = path.find_last_of("/\\");
48  if (tmp == path.length() - 1) {
49  size_t i = tmp;
50  while ((path[i] == '/' || path[i] == '\\') && i >= 0) {
51  --i;
52  }
53  path_ = path.substr(0, i + 1);
54  } else {
55  path_ = path;
56  }
57  std::vector<char> fname(path_.length() + 1);
58  std::vector<char> ext(path_.length() + 1);
59  _splitpath_s(path_.c_str(), NULL, 0, NULL, 0,
60  &fname[0], path_.length() + 1, &ext[0], path_.length() + 1);
61  return std::string(&fname[0]) + std::string(&ext[0]);
62 #else
63  char* path_ = strdup(path.c_str());
64  char* base = basename(path_);
65  std::string ret(base);
66  free(path_);
67  return ret;
68 #endif
69 }
70 
71 inline void HandleSystemError(const std::string& msg) {
72 #ifdef _WIN32
73  LPVOID msg_buf;
74  DWORD dw = GetLastError();
75  FormatMessage(
76  FORMAT_MESSAGE_ALLOCATE_BUFFER |
77  FORMAT_MESSAGE_FROM_SYSTEM |
78  FORMAT_MESSAGE_IGNORE_INSERTS,
79  NULL,
80  dw,
81  MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
82  (LPTSTR)&msg_buf,
83  0, NULL);
84  const std::string msg_err(static_cast<const char*>(msg_buf));
85  LocalFree(msg_buf);
86 #else
87  const std::string msg_err(strerror(errno));
88 #endif
89  LOG(FATAL) << msg << "\nReason: " << msg_err;
90 }
91 
92 inline void CreateDirectoryIfNotExist(const char* dirpath) {
93 #ifdef _WIN32
94  DWORD ftyp = GetFileAttributesA(dirpath);
95  if (ftyp == INVALID_FILE_ATTRIBUTES) {
96  // directory doesn't seem to exist; attempt to create one
97  if (CreateDirectoryA(dirpath, NULL) == 0) {
98  // failed to create a new directory
99  HandleSystemError(std::string("CreateDirectoryIfNotExist: "
100  "failed to create new directory ") + dirpath);
101  }
102  } else {
103  if (!(ftyp & FILE_ATTRIBUTE_DIRECTORY)) {
104  LOG(FATAL) << "CreateDirectoryIfNotExist: "
105  << dirpath << " is a file, not a directory";
106  }
107  }
108 #else
109  struct stat sb;
110  if (stat(dirpath, &sb) != 0) {
111  // directory doesn't seem to exist; attempt to create one
112  if (mkdir(dirpath, S_IRUSR | S_IWUSR | S_IXUSR) != 0) {
113  // failed to create a new directory
114  HandleSystemError(std::string("CreateDirectoryIfNotExist: "
115  "failed to create new directory ") + dirpath);
116  }
117  } else {
118  if (!S_ISDIR(sb.st_mode)) {
119  LOG(FATAL) << "CreateDirectoryIfNotExist: "
120  << dirpath << " is a file, not a directory";
121  }
122  }
123 #endif
124 }
125 
126 inline void CreateDirectoryIfNotExistRecursive(const std::string& dirpath) {
127  std::string dirpath_;
128 #ifdef _WIN32
129  if (dirpath.find("/") == std::string::npos
130  && dirpath.find("\\") != std::string::npos) {
131  // replace backward slashes with forward slashes
132  dirpath_ = std::regex_replace(dirpath, std::regex("\\\\"), "/");
133  } else {
134  dirpath_ = dirpath;
135  }
136 #else
137  dirpath_ = dirpath;
138 #endif
139  const std::vector<std::string> tokens = common::Split(dirpath_, '/');
140  std::string accum;
141  size_t i;
142  if (tokens[0].empty()) { // absolute path, starting with '/'
143  accum = "/" + tokens[1];
144  i = 1;
145  } else { // relative path
146  accum = tokens[0];
147  i = 0;
148  }
149  for (; i < tokens.size(); ++i) {
150  common::filesystem::CreateDirectoryIfNotExist(accum.c_str());
151  if (i < tokens.size() - 1 && !tokens[i + 1].empty()) {
152  accum += "/";
153  accum += tokens[i + 1];
154  }
155  }
156 }
157 
159  public:
161 #if _WIN32
162  /* locate the root directory of temporary area */
163  char tmproot[MAX_PATH] = {0};
164  const DWORD dw_retval = GetTempPathA(MAX_PATH, tmproot);
165  if (dw_retval > MAX_PATH || dw_retval == 0) {
166  LOG(FATAL) << "TemporaryDirectory(): "
167  << "Could not create temporary directory";
168  }
169  /* generate a unique 8-letter alphanumeric string */
170  const std::string letters = "abcdefghijklmnopqrstuvwxyz0123456789_";
171  std::string uniqstr(8, '\0');
172  std::random_device rd;
173  std::mt19937 gen(rd());
174  std::uniform_int_distribution<int> dis(0, letters.length() - 1);
175  std::generate(uniqstr.begin(), uniqstr.end(),
176  [&dis, &gen, &letters]() -> char {
177  return letters[dis(gen)];
178  });
179  /* combine paths to get the name of the temporary directory */
180  char tmpdir[MAX_PATH] = {0};
181  PathCombineA(tmpdir, tmproot, uniqstr.c_str());
182  if (!CreateDirectoryA(tmpdir, NULL)) {
183  LOG(FATAL) << "TemporaryDirectory(): "
184  << "Could not create temporary directory";
185  }
186  path = std::string(tmpdir);
187 #else
188  std::string tmproot; /* root directory of temporary area */
189  std::string dirtemplate; /* template for temporary directory name */
190  /* Get TMPDIR env variable or fall back to /tmp/ */
191  {
192  const char* tmpenv = getenv("TMPDIR");
193  if (tmpenv) {
194  tmproot = std::string(tmpenv);
195  // strip trailing forward slashes
196  while (tmproot.length() != 0 && tmproot[tmproot.length() - 1] == '/') {
197  tmproot.resize(tmproot.length() - 1);
198  }
199  } else {
200  tmproot = "/tmp";
201  }
202  }
203  dirtemplate = tmproot + "/tmpdir.XXXXXX";
204  std::vector<char> dirtemplate_buf(dirtemplate.begin(), dirtemplate.end());
205  dirtemplate_buf.push_back('\0');
206  char* tmpdir = mkdtemp(&dirtemplate_buf[0]);
207  if (!tmpdir) {
208  LOG(FATAL) << "TemporaryDirectory(): "
209  << "Could not create temporary directory";
210  }
211  path = std::string(tmpdir);
212 #endif
213  LOG(INFO) << "Created temporary directory " << path;
214  }
215  ~TemporaryDirectory() {
216  for (const std::string& filename : file_list) {
217  if (std::remove(filename.c_str()) != 0) {
218  LOG(FATAL) << "Couldn't remove file " << filename;
219  }
220  }
221 #if _WIN32
222  const bool rmdir_success = (RemoveDirectoryA(path.c_str()) != 0);
223 #else
224  const bool rmdir_success = (rmdir(path.c_str()) == 0);
225 #endif
226  if (rmdir_success) {
227  LOG(INFO) << "Successfully deleted temporary directory " << path;
228  } else {
229  LOG(FATAL) << "~TemporaryDirectory(): "
230  << "Could not remove temporary directory ";
231  }
232  }
233 
234  std::string AddFile(const std::string& filename) {
235  const std::string file_path = this->path + "/" + filename;
236  file_list.push_back(file_path);
237  return file_path;
238  }
239 
240  std::string path;
241 
242  private:
243  std::vector<std::string> file_list;
244 };
245 
246 } // namespace filesystem
247 } // namespace common
248 } // namespace treelite
249 
250 #endif // TREELITE_COMMON_FILESYSTEM_H_