xref: /aosp_15_r20/external/tensorflow/tensorflow/python/tools/module_util.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Helper functions for modules."""
16import os
17
18import importlib
19
20
21def get_parent_dir(module):
22  return os.path.abspath(os.path.join(os.path.dirname(module.__file__), ".."))
23
24
25def get_parent_dir_for_name(module_name):
26  """Get parent directory for module with the given name.
27
28  Args:
29    module_name: Module name for e.g.
30      tensorflow_estimator.python.estimator.api._v1.estimator.
31
32  Returns:
33    Path to the parent directory if module is found and None otherwise.
34    Given example above, it should return:
35      /pathtoestimator/tensorflow_estimator/python/estimator/api/_v1.
36  """
37  name_split = module_name.split(".")
38  if not name_split:
39    return None
40
41  try:
42    spec = importlib.util.find_spec(name_split[0])
43  except ValueError:
44    return None
45  if not spec or not spec.origin:
46    return None
47  base_path = os.path.dirname(spec.origin)
48  return os.path.join(base_path, *name_split[1:-1])
49