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