[TF-DF] Add migration message for OSS

PiperOrigin-RevId: 690588920

tensorflow_decision_forests/__init__.py

@@ -70,24 +70,38 @@
 
 import os
 import sys
+import io
+
+
+def _is_direct_output(stream=sys.stdout):
+  """Checks if output stream redirects to the shell/console directly."""
+
+  if stream.isatty():
+    return True
+  if isinstance(stream, io.TextIOWrapper):
+    return _is_direct_output(stream.buffer)
+  if isinstance(stream, io.BufferedWriter):
+    return _is_direct_output(stream.raw)
+  if isinstance(stream, io.FileIO):
+    return stream.fileno() in [1, 2]
+  return False
+
 
 # Only print the welcome message if TFDF_DISABLE_WELCOME_MESSAGE is not set and
-# if user has notalready imported YDF.
+# if user has not already imported YDF.
 if (
     os.getenv("TFDF_DISABLE_WELCOME_MESSAGE") is None
     and "ydf" not in sys.modules
 ):
 
-  if (
-      "IPython" in sys.modules and "google.colab" in sys.modules
-  ):  # Check if executed in a Notebook
+  if not _is_direct_output():  # Check if executed in a Notebook
 
     import IPython
 
     IPython.display.display(IPython.display.HTML("""
 <p style="margin:0px;">🌲 Try <a href="https://ydf.readthedocs.io/en/latest/" target="_blank">YDF</a>, the successor of
     <a href="https://www.tensorflow.org/decision_forests" target="_blank">TensorFlow
-        Decision Forests</a> using the same algorithm but with more features and faster
+        Decision Forests</a> using the same algorithms but with more features and faster
     training!
 </p>
 <div style="display: flex; flex-wrap: wrap; margin:5px;max-width: 880px;">
@@ -121,4 +135,16 @@
 """))
 
   else:
+    import termcolor
+    print(
+        termcolor.colored("🌲 Try ", "green"),
+        termcolor.colored("https://ydf.readthedocs.io", "blue"),
+        termcolor.colored(
+            ", the successor of TensorFlow Decision Forests with more"
+            " features and faster training!",
+            "green",
+        ),
+        sep="",
+        file=sys.stderr,
+    )
     pass