Skip to content

Commit c3f2df3

Browse files
rstzcopybara-github
authored andcommitted
[TF-DF] Add migration message for OSS
PiperOrigin-RevId: 690588920
1 parent 41a8f56 commit c3f2df3

File tree

1 file changed

+31
-5
lines changed

1 file changed

+31
-5
lines changed

tensorflow_decision_forests/__init__.py

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -70,24 +70,38 @@
7070

7171
import os
7272
import sys
73+
import io
74+
75+
76+
def _is_direct_output(stream=sys.stdout):
77+
"""Checks if output stream redirects to the shell/console directly."""
78+
79+
if stream.isatty():
80+
return True
81+
if isinstance(stream, io.TextIOWrapper):
82+
return _is_direct_output(stream.buffer)
83+
if isinstance(stream, io.BufferedWriter):
84+
return _is_direct_output(stream.raw)
85+
if isinstance(stream, io.FileIO):
86+
return stream.fileno() in [1, 2]
87+
return False
88+
7389

7490
# Only print the welcome message if TFDF_DISABLE_WELCOME_MESSAGE is not set and
75-
# if user has notalready imported YDF.
91+
# if user has not already imported YDF.
7692
if (
7793
os.getenv("TFDF_DISABLE_WELCOME_MESSAGE") is None
7894
and "ydf" not in sys.modules
7995
):
8096

81-
if (
82-
"IPython" in sys.modules and "google.colab" in sys.modules
83-
): # Check if executed in a Notebook
97+
if not _is_direct_output(): # Check if executed in a Notebook
8498

8599
import IPython
86100

87101
IPython.display.display(IPython.display.HTML("""
88102
<p style="margin:0px;">🌲 Try <a href="https://ydf.readthedocs.io/en/latest/" target="_blank">YDF</a>, the successor of
89103
<a href="https://www.tensorflow.org/decision_forests" target="_blank">TensorFlow
90-
Decision Forests</a> using the same algorithm but with more features and faster
104+
Decision Forests</a> using the same algorithms but with more features and faster
91105
training!
92106
</p>
93107
<div style="display: flex; flex-wrap: wrap; margin:5px;max-width: 880px;">
@@ -121,4 +135,16 @@
121135
"""))
122136

123137
else:
138+
import termcolor
139+
print(
140+
termcolor.colored("🌲 Try ", "green"),
141+
termcolor.colored("https://ydf.readthedocs.io", "blue"),
142+
termcolor.colored(
143+
", the successor of TensorFlow Decision Forests with more"
144+
" features and faster training!",
145+
"green",
146+
),
147+
sep="",
148+
file=sys.stderr,
149+
)
124150
pass

0 commit comments

Comments
 (0)